package auth

import (
	"context"
	"crypto"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"math/big"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"
)

type contextKey string

// UserKey is the context key for authenticated user claims.
const UserKey contextKey = "user"

// Audience handles either a single string or an array of strings in JWT audience claim.
type Audience []string

func (a *Audience) UnmarshalJSON(b []byte) error {
	if len(b) > 0 && b[0] == '"' {
		var s string
		if err := json.Unmarshal(b, &s); err != nil {
			return err
		}
		*a = Audience{s}
		return nil
	}
	return json.Unmarshal(b, (*[]string)(a))
}

func (a Audience) Contains(s string) bool {
	for _, v := range a {
		if v == s {
			return true
		}
	}
	return false
}

// IDClaims represents verified claims in an OIDC token.
type IDClaims struct {
	Iss   string   `json:"iss"`
	Sub   string   `json:"sub"`
	Aud   Audience `json:"aud"`
	Email string   `json:"email,omitempty"`
	Name  string   `json:"name,omitempty"`
	Roles []string `json:"roles,omitempty"`
	Scope string   `json:"scope,omitempty"`
	Exp   int64    `json:"exp"`
}

type jwksDoc struct {
	Keys []struct {
		Kid string `json:"kid"`
		Kty string `json:"kty"`
		Alg string `json:"alg"`
		N   string `json:"n"`
		E   string `json:"e"`
	} `json:"keys"`
}

// OidcAuth provides OIDC token verification and authentication middleware.
type OidcAuth struct {
	issuer     string
	client     string
	disabled   bool
	ingestKey  string
	agentToken string
	// audience is an optional extra accepted token audience: this service's
	// own URL, which the issuer emits via AuthClient.spec.audience (e.g. for
	// device-grant dev tokens minted for another client).
	audience string

	mu      sync.Mutex
	keys    map[string]*rsa.PublicKey
	fetched time.Time
}

// NewOidcAuth constructs a new OidcAuth instance.
func NewOidcAuth(ingestKey string) *OidcAuth {
	issuer := os.Getenv("OTEL_OIDC_ISSUER")
	if issuer == "" {
		issuer = "https://users.k6n.net"
	}
	client := os.Getenv("OTEL_OIDC_CLIENT_ID")
	if client == "" {
		client = "otel-debug"
	}
	disabled := os.Getenv("OTEL_AUTH_DISABLED") == "true"

	return &OidcAuth{
		issuer:     strings.TrimSuffix(issuer, "/"),
		client:     client,
		disabled:   disabled,
		ingestKey:  ingestKey,
		agentToken: os.Getenv("OTEL_AGENT_TOKEN"),
		audience:   os.Getenv("OTEL_OIDC_AUDIENCE"),
	}
}

func (a *OidcAuth) Issuer() string {
	return a.issuer
}

func (a *OidcAuth) ClientID() string {
	return a.client
}

func (a *OidcAuth) IsDisabled() bool {
	return a.disabled
}

func (a *OidcAuth) fetchKeys(ctx context.Context) (map[string]*rsa.PublicKey, error) {
	url := a.issuer + "/oidc/v1/keys"
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("jwks endpoint returned %d", resp.StatusCode)
	}

	var doc jwksDoc
	if err := json.NewDecoder(resp.Body).Decode(&doc); err != nil {
		return nil, err
	}

	keys := make(map[string]*rsa.PublicKey, len(doc.Keys))
	for _, k := range doc.Keys {
		if k.Kty != "RSA" {
			continue
		}
		n, err := base64.RawURLEncoding.DecodeString(k.N)
		if err != nil {
			continue
		}
		eBytes, err := base64.RawURLEncoding.DecodeString(k.E)
		if err != nil || len(eBytes) > 4 {
			continue
		}
		var eInt uint32
		for _, b := range eBytes {
			eInt = eInt<<8 | uint32(b)
		}
		keys[k.Kid] = &rsa.PublicKey{N: new(big.Int).SetBytes(n), E: int(eInt)}
	}

	if len(keys) == 0 {
		return nil, errors.New("jwks contained no RSA keys")
	}
	return keys, nil
}

func (a *OidcAuth) publicKey(ctx context.Context, kid string) (*rsa.PublicKey, error) {
	a.mu.Lock()
	defer a.mu.Unlock()

	if key, ok := a.keys[kid]; ok {
		return key, nil
	}

	// Throttle refreshes to at most once per minute
	if a.keys == nil || time.Since(a.fetched) > time.Minute {
		keys, err := a.fetchKeys(ctx)
		if err != nil {
			return nil, err
		}
		a.keys = keys
		a.fetched = time.Now()
	}

	key, ok := a.keys[kid]
	if !ok {
		return nil, fmt.Errorf("unknown key id %q", kid)
	}
	return key, nil
}

// VerifyToken verifies an RS256 JWT against the cached JWKS public keys.
func (a *OidcAuth) VerifyToken(ctx context.Context, token string) (*IDClaims, error) {
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return nil, errors.New("malformed token")
	}

	var header struct {
		Kid string `json:"kid"`
		Alg string `json:"alg"`
	}
	h, err := base64.RawURLEncoding.DecodeString(parts[0])
	if err != nil {
		return nil, fmt.Errorf("decode header: %w", err)
	}
	if err := json.Unmarshal(h, &header); err != nil {
		return nil, fmt.Errorf("parse header: %w", err)
	}
	if header.Alg != "RS256" {
		return nil, fmt.Errorf("unexpected alg %q", header.Alg)
	}

	key, err := a.publicKey(ctx, header.Kid)
	if err != nil {
		return nil, err
	}

	sig, err := base64.RawURLEncoding.DecodeString(parts[2])
	if err != nil {
		return nil, fmt.Errorf("decode signature: %w", err)
	}

	sum := sha256.Sum256([]byte(parts[0] + "." + parts[1]))
	if err := rsa.VerifyPKCS1v15(key, crypto.SHA256, sum[:], sig); err != nil {
		return nil, errors.New("invalid signature")
	}

	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return nil, fmt.Errorf("decode claims: %w", err)
	}

	var claims IDClaims
	if err := json.Unmarshal(payload, &claims); err != nil {
		return nil, fmt.Errorf("parse claims: %w", err)
	}

	if claims.Iss != a.issuer {
		return nil, fmt.Errorf("unexpected issuer %q", claims.Iss)
	}
	if !claims.Aud.Contains(a.client) && (a.audience == "" || !claims.Aud.Contains(a.audience)) {
		return nil, errors.New("token not issued for this client")
	}
	if claims.Exp != 0 && time.Now().Unix() > claims.Exp {
		return nil, errors.New("token expired")
	}

	return &claims, nil
}

// HandleTokenExchange proxies the OIDC token exchange endpoint so the SPA
// doesn't face CORS issues.
func (a *OidcAuth) HandleTokenExchange(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "POST required", http.StatusMethodNotAllowed)
		return
	}
	if err := r.ParseForm(); err != nil {
		http.Error(w, "invalid form body", http.StatusBadRequest)
		return
	}

	// Always ensure client_id is set
	if r.PostForm.Get("client_id") == "" {
		r.PostForm.Set("client_id", a.client)
	}

	resp, err := http.PostForm(a.issuer+"/oauth/v2/token", r.PostForm)
	if err != nil {
		log.Printf("token endpoint unreachable: %v", err)
		http.Error(w, "token endpoint unreachable", http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(resp.StatusCode)
	_, _ = io.Copy(w, resp.Body)
}

// HandleConfig returns OIDC client configuration for the SPA frontend.
func (a *OidcAuth) HandleConfig(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(map[string]interface{}{
		"issuer":   a.issuer,
		"clientId": a.client,
		"disabled": a.disabled,
	})
}

// RequireAuth wraps an HTTP handler and ensures valid authentication.
func (a *OidcAuth) RequireAuth(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if a.disabled {
			next.ServeHTTP(w, r)
			return
		}

		// 1. Allow if upstream proxy (ingress) already verified via auth-url
		if authUser := r.Header.Get("X-Auth-User"); authUser != "" {
			ctx := context.WithValue(r.Context(), UserKey, &IDClaims{Sub: authUser, Email: authUser})
			next.ServeHTTP(w, r.WithContext(ctx))
			return
		}

		// 2. Allow if valid ingest key is provided (for background automation / agents)
		if a.ingestKey != "" {
			k := r.Header.Get("X-Otel-Ingest-Key")
			if k == "" {
				k = r.URL.Query().Get("key")
			}
			if k != "" && k == a.ingestKey {
				next.ServeHTTP(w, r)
				return
			}
		}

		// Service-to-service credential for debugging agents. It is separate
		// from the ingest key because it grants access to protected debug data.
		if a.agentToken != "" && constantTimeTokenMatch(r, a.agentToken) {
			ctx := context.WithValue(r.Context(), UserKey, &IDClaims{Sub: "service-agent", Name: "debug agent"})
			next.ServeHTTP(w, r.WithContext(ctx))
			return
		}

		// 3. Check Bearer token (or access_token query param for WebSocket)
		token := bearerToken(r)
		if token != "" {
			claims, err := a.VerifyToken(r.Context(), token)
			if err == nil {
				ctx := context.WithValue(r.Context(), UserKey, claims)
				next.ServeHTTP(w, r.WithContext(ctx))
				return
			}
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusUnauthorized)
			_ = json.NewEncoder(w).Encode(map[string]string{
				"error": "invalid token: " + err.Error(),
			})
			return
		}

		// 4. Unauthorized
		w.Header().Set("WWW-Authenticate", fmt.Sprintf(`Bearer realm="%s"`, a.client))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusUnauthorized)
		_ = json.NewEncoder(w).Encode(map[string]string{
			"error": "authentication required",
		})
	})
}

func constantTimeTokenMatch(r *http.Request, expected string) bool {
	token := r.Header.Get("X-Otel-Agent-Token")
	if token == "" {
		auth := r.Header.Get("Authorization")
		if len(auth) > 7 && strings.EqualFold(auth[:7], "bearer ") {
			token = strings.TrimSpace(auth[7:])
		}
	}
	return len(token) == len(expected) && subtle.ConstantTimeCompare([]byte(token), []byte(expected)) == 1
}

func bearerToken(r *http.Request) string {
	if t := r.URL.Query().Get("access_token"); t != "" {
		return t
	}
	auth := r.Header.Get("Authorization")
	if len(auth) > 7 && strings.EqualFold(auth[:7], "bearer ") {
		return strings.TrimSpace(auth[7:])
	}
	return ""
}

// CurrentUser retrieves the authenticated user claims from the context, if present.
func CurrentUser(ctx context.Context) *IDClaims {
	claims, _ := ctx.Value(UserKey).(*IDClaims)
	return claims
}