package main

import (
	"context"
	"crypto/rsa"
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"math/big"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/golang-jwt/jwt/v5"
)

type contextKey string

const userContextKey contextKey = "auth_user"

type User struct {
	Sub   string
	Email string
	Name  string
}

type Auth struct {
	issuer   string
	audience string
	jwksURL  string
	devMode  bool
	db       *sql.DB

	mu         sync.RWMutex
	keyByKid   map[string]*rsa.PublicKey
	lastSyncAt time.Time
}

func NewAuth(cfg Config) (*Auth, error) {
	a := &Auth{
		issuer:    cfg.OIDCIssuer,
		audience:  cfg.OIDCAudience,
		jwksURL:   cfg.OIDCJWKSURL,
		devMode:   cfg.DevMode,
		keyByKid:  map[string]*rsa.PublicKey{},
		lastSyncAt: time.Time{},
	}
	if a.devMode {
		return a, nil
	}
	if err := a.refreshKeys(context.Background()); err != nil {
		return nil, err
	}
	return a, nil
}

func (a *Auth) WithDB(db *sql.DB) {
	a.db = db
}

func (a *Auth) RequireAuth(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var (
			user User
			err  error
		)
		if a.devMode {
			user = User{
				Sub:   "demo-user-001",
				Email: "demo@example.com",
				Name:  "Demo User",
			}
		} else {
			user, err = a.authenticateRequest(r)
			if err != nil {
				http.Error(w, "unauthorized", http.StatusUnauthorized)
				return
			}
		}

		if a.db != nil {
			_, _ = a.db.ExecContext(r.Context(), `
				INSERT INTO users (sub, email, name)
				VALUES ($1, $2, $3)
				ON CONFLICT (sub) DO UPDATE
				SET email = EXCLUDED.email, name = EXCLUDED.name
			`, user.Sub, user.Email, user.Name)
		}

		ctx := context.WithValue(r.Context(), userContextKey, user)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

func MustUserFromContext(ctx context.Context) User {
	v := ctx.Value(userContextKey)
	if v == nil {
		return User{}
	}
	if u, ok := v.(User); ok {
		return u
	}
	return User{}
}

func (a *Auth) authenticateRequest(r *http.Request) (User, error) {
	token := extractBearer(r.Header.Get("Authorization"))
	if token == "" {
		return User{}, errors.New("missing token")
	}

	parser := jwt.NewParser(jwt.WithValidMethods([]string{"RS256"}))
	parsed, err := parser.Parse(token, func(t *jwt.Token) (any, error) {
		kid, _ := t.Header["kid"].(string)
		if kid == "" {
			return nil, errors.New("missing kid")
		}
		key, err := a.lookupKey(r.Context(), kid)
		if err != nil {
			return nil, err
		}
		return key, nil
	})
	if err != nil || !parsed.Valid {
		return User{}, errors.New("invalid token")
	}

	claims, ok := parsed.Claims.(jwt.MapClaims)
	if !ok {
		return User{}, errors.New("invalid claims")
	}

	if iss, _ := claims["iss"].(string); iss != a.issuer {
		return User{}, errors.New("invalid issuer")
	}
	if !audienceMatches(claims["aud"], a.audience) {
		return User{}, errors.New("invalid audience")
	}

	sub, _ := claims["sub"].(string)
	email, _ := claims["email"].(string)
	name, _ := claims["name"].(string)
	if sub == "" {
		return User{}, errors.New("missing sub")
	}
	return User{Sub: sub, Email: email, Name: name}, nil
}

func (a *Auth) lookupKey(ctx context.Context, kid string) (*rsa.PublicKey, error) {
	a.mu.RLock()
	key := a.keyByKid[kid]
	last := a.lastSyncAt
	a.mu.RUnlock()
	if key != nil {
		return key, nil
	}
	if time.Since(last) > time.Minute {
		if err := a.refreshKeys(ctx); err != nil {
			return nil, err
		}
		a.mu.RLock()
		defer a.mu.RUnlock()
		if refreshed := a.keyByKid[kid]; refreshed != nil {
			return refreshed, nil
		}
	}
	return nil, fmt.Errorf("no key for kid %s", kid)
}

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

func (a *Auth) refreshKeys(ctx context.Context) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, a.jwksURL, nil)
	if err != nil {
		return err
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("jwks request failed: %s", resp.Status)
	}

	var doc jwksDoc
	if err := json.NewDecoder(resp.Body).Decode(&doc); err != nil {
		return err
	}
	newKeys := map[string]*rsa.PublicKey{}
	for _, k := range doc.Keys {
		if k.Kty != "RSA" || k.Kid == "" || k.N == "" || k.E == "" {
			continue
		}
		pub, err := rsaKeyFromJWK(k.N, k.E)
		if err != nil {
			continue
		}
		newKeys[k.Kid] = pub
	}
	if len(newKeys) == 0 {
		return errors.New("no usable jwks keys")
	}

	a.mu.Lock()
	a.keyByKid = newKeys
	a.lastSyncAt = time.Now()
	a.mu.Unlock()
	return nil
}

func rsaKeyFromJWK(nB64, eB64 string) (*rsa.PublicKey, error) {
	nBytes, err := base64.RawURLEncoding.DecodeString(nB64)
	if err != nil {
		return nil, err
	}
	eBytes, err := base64.RawURLEncoding.DecodeString(eB64)
	if err != nil {
		return nil, err
	}
	n := new(big.Int).SetBytes(nBytes)
	e := 0
	for _, b := range eBytes {
		e = e<<8 + int(b)
	}
	if e == 0 {
		return nil, errors.New("invalid exponent")
	}
	return &rsa.PublicKey{N: n, E: e}, nil
}

func extractBearer(value string) string {
	if value == "" {
		return ""
	}
	parts := strings.SplitN(value, " ", 2)
	if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") {
		return ""
	}
	return strings.TrimSpace(parts[1])
}

func audienceMatches(audValue any, expected string) bool {
	switch v := audValue.(type) {
	case string:
		return v == expected
	case []any:
		for _, item := range v {
			s, ok := item.(string)
			if ok && s == expected {
				return true
			}
		}
	}
	return false
}