package jwtauth

import (
	"context"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"strings"
	"time"

	"github.com/go-jose/go-jose/v4"
	"github.com/go-jose/go-jose/v4/jwt"

	"commilitia.net/cdrop/internal/config"
	"commilitia.net/cdrop/internal/db"
)

// Store is the subset of *db.Queries the auth middleware uses: device upserts
// plus the shortcut-token lookups needed to honour revocation. Declared as an
// interface so tests can swap in a fake.
type Store interface {
	UpsertDevice(ctx context.Context, arg db.UpsertDeviceParams) error
	GetShortcutToken(ctx context.Context, jti string) (db.ShortcutToken, error)
	TouchShortcutTokenUsed(ctx context.Context, arg db.TouchShortcutTokenUsedParams) error
}

type Authenticator struct {
	cfg             *config.Config
	store           Store
	jwks            *jwksCache
	hsKey           []byte
	sessionTokenKey []byte
}

func New(cfg *config.Config, store Store) *Authenticator {
	a := &Authenticator{
		cfg:   cfg,
		store: store,
	}
	if cfg.HS256Secret != "" {
		a.hsKey = DeriveHS256Key(cfg.HS256Secret)
	}
	// Self-signed session tokens (scan-login) are keyed off SessionSecret; nil
	// when unset (dev) disables verifySelfToken, matching the minting side.
	a.sessionTokenKey = DeriveSessionTokenKey(cfg.SessionSecret)
	if cfg.AuthMode == "prod" && cfg.OIDCJWKSURL != "" {
		a.jwks = newJWKSCache(cfg.OIDCJWKSURL, 10*time.Minute)
	}
	return a
}

func (a *Authenticator) Middleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		token, ok := bearerToken(r)
		if !ok {
			unauthorized(w, "missing bearer token")
			return
		}

		claims, err := a.verify(r.Context(), token, r)
		if err != nil {
			slog.Warn("auth failed", "err", err, "path", r.URL.Path)
			unauthorized(w, "invalid token")
			return
		}

		deviceName := SanitizeDeviceName(r.Header.Get("X-Device-Name"))
		deviceType := normalizeDeviceType(r.Header.Get("X-Device-Type"))
		if claims.Scoped() {
			// The backend authoritatively knows a scoped token is the iOS
			// Shortcut, so it labels the device "shortcut" regardless of any
			// header. ("ios" stays reserved for a real native client.)
			deviceType = "shortcut"
		}

		// Register a device only when the client names itself (X-Device-Name).
		// A nameless request — e.g. a polling shortcut hitting /api/clipboard/version
		// without the header — must NOT be registered: the old random UA fallback
		// minted a fresh "Unknown Device" row on every request and flooded the list.
		if deviceName != "" {
			if err := a.store.UpsertDevice(r.Context(), db.UpsertDeviceParams{
				UserID:   claims.UserID,
				Name:     deviceName,
				Type:     deviceType,
				LastSeen: time.Now().Unix(),
			}); err != nil {
				// non-fatal: log and continue so transient DB errors don't 401 users
				slog.Error("device upsert failed",
					"err", err, "user", claims.UserID, "device", deviceName)
			}
		}

		ctx := context.WithValue(r.Context(), claimsCtxKey, claims)
		ctx = context.WithValue(ctx, deviceCtxKey, deviceName)
		ctx = context.WithValue(ctx, deviceTypeCtxKey, deviceType)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

func (a *Authenticator) verify(ctx context.Context, token string, r *http.Request) (*Claims, error) {
	if a.cfg.AuthMode == "dev" {
		return a.verifyDev(token, r)
	}
	// cdrop self-signed session tokens first: distinct key + typ=session, so a
	// shortcut token (different key, requires jti) never validates here and a
	// session token never falls through to the shortcut path's DB lookup.
	if c, err := a.verifySelfToken(token); err == nil {
		return c, nil
	}
	if c, err := a.verifyHS256(ctx, token); err == nil {
		return c, nil
	}
	return a.verifyRS256(ctx, token)
}

// verifySelfToken validates a cdrop self-signed session access token (AUTH.md
// §3.1): HS256 over DeriveSessionTokenKey, carrying typ=session and a full/guest
// scope. Stateless by design — no DB lookup, so it stays cheap on every request;
// revocation rides the short TTL (the session row is re-checked at /auth/refresh).
func (a *Authenticator) verifySelfToken(token string) (*Claims, error) {
	if len(a.sessionTokenKey) == 0 {
		return nil, errors.New("session token key not configured")
	}
	parsed, err := jwt.ParseSigned(token, []jose.SignatureAlgorithm{jose.HS256})
	if err != nil {
		return nil, err
	}
	var std jwt.Claims
	custom := map[string]any{}
	if err := parsed.Claims(a.sessionTokenKey, &std, &custom); err != nil {
		return nil, err
	}
	if err := std.ValidateWithLeeway(jwt.Expected{Time: time.Now()}, 30*time.Second); err != nil {
		return nil, err
	}
	if typ, _ := custom["typ"].(string); typ != "session" {
		return nil, errors.New("not a session token")
	}
	if std.Subject == "" {
		return nil, errors.New("session token missing subject")
	}
	scope, _ := custom["scope"].(string)
	if scope != "full" && scope != "guest" {
		return nil, errors.New("session token invalid scope")
	}
	return &Claims{UserID: std.Subject, SessionScope: scope}, nil
}

func (a *Authenticator) verifyDev(token string, r *http.Request) (*Claims, error) {
	if subtle.ConstantTimeCompare([]byte(token), []byte(a.cfg.DevToken)) != 1 {
		return nil, errors.New("invalid dev token")
	}
	userID := r.Header.Get("X-Dev-User")
	if userID == "" {
		userID = "dev-user"
	}
	return &Claims{UserID: userID, Groups: []string{"dev"}}, nil
}

func (a *Authenticator) verifyHS256(ctx context.Context, token string) (*Claims, error) {
	if len(a.hsKey) == 0 {
		return nil, errors.New("HS256 secret not configured")
	}
	parsed, err := jwt.ParseSigned(token, []jose.SignatureAlgorithm{jose.HS256})
	if err != nil {
		return nil, err
	}
	var std jwt.Claims
	custom := map[string]any{}
	if err := parsed.Claims(a.hsKey, &std, &custom); err != nil {
		return nil, err
	}
	if err := std.ValidateWithLeeway(jwt.Expected{Time: time.Now()}, 30*time.Second); err != nil {
		return nil, err
	}
	claims, err := claimsFromJWT(std, custom)
	if err != nil {
		return nil, err
	}

	// HS256 is only ever used to mint scoped shortcut tokens, and those ALWAYS
	// carry a jti. A validly-signed HS256 token without one must be rejected
	// outright: otherwise it would fall through here as a full, unscoped account
	// session with a self-declared subject — far beyond the clipboard-only
	// surface HS256 is meant for. Requiring the jti keeps a leaked HS256 secret's
	// blast radius pinned to the shortcut scope (clipboard, and nothing else).
	if std.ID == "" {
		return nil, errors.New("HS256 token missing jti")
	}

	// The signature alone is not enough — the token must still be present and not
	// revoked in the store, so a leaked or retired token can be killed
	// server-side. The stored row is also the authoritative source of scopes
	// (never trust scopes off the wire).
	row, err := a.store.GetShortcutToken(ctx, std.ID)
	if err != nil {
		return nil, fmt.Errorf("shortcut token lookup: %w", err)
	}
	if row.Revoked != 0 {
		return nil, errors.New("shortcut token revoked")
	}
	if row.UserID != claims.UserID {
		return nil, errors.New("shortcut token subject mismatch")
	}
	claims.JTI = std.ID
	claims.Scopes = strings.Fields(row.Scopes)
	a.touchTokenAsync(std.ID)
	return claims, nil
}

// touchTokenAsync records a shortcut token's last-use time off the request path
// — it's audit metadata, so a slow or failing write must never delay or fail
// the request it belongs to.
func (a *Authenticator) touchTokenAsync(jti string) {
	now := time.Now().Unix()
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		if err := a.store.TouchShortcutTokenUsed(ctx, db.TouchShortcutTokenUsedParams{
			LastUsedAt: &now,
			Jti:        jti,
		}); err != nil {
			slog.Warn("touch shortcut token failed", "err", err, "jti", jti)
		}
	}()
}

func (a *Authenticator) verifyRS256(ctx context.Context, token string) (*Claims, error) {
	if a.jwks == nil {
		return nil, errors.New("JWKS not configured")
	}
	parsed, err := jwt.ParseSigned(token, []jose.SignatureAlgorithm{jose.RS256})
	if err != nil {
		return nil, err
	}
	if len(parsed.Headers) == 0 {
		return nil, errors.New("missing token headers")
	}
	kid := parsed.Headers[0].KeyID
	key, err := a.jwks.GetKey(ctx, kid)
	if err != nil {
		return nil, err
	}
	var std jwt.Claims
	custom := map[string]any{}
	if err := parsed.Claims(key, &std, &custom); err != nil {
		return nil, err
	}
	expected := jwt.Expected{Time: time.Now()}
	if a.cfg.OIDCIssuer != "" {
		expected.Issuer = a.cfg.OIDCIssuer
	}
	if a.cfg.OIDCAudience != "" {
		expected.AnyAudience = parseAudiences(a.cfg.OIDCAudience)
	}
	if err := std.ValidateWithLeeway(expected, 30*time.Second); err != nil {
		return nil, err
	}
	return claimsFromJWT(std, custom)
}

// VerifyIDToken validates an OIDC id_token (RS256 via JWKS, issuer + audience)
// from a fresh prompt=login exchange and returns its subject and auth_time (the
// epoch second of the actual end-user authentication). Used for step-up re-auth
// (AUTH.md §6): the caller checks sub matches and auth_time is recent enough.
func (a *Authenticator) VerifyIDToken(ctx context.Context, token string) (subject string, authTime int64, err error) {
	if a.jwks == nil {
		return "", 0, errors.New("JWKS not configured")
	}
	parsed, err := jwt.ParseSigned(token, []jose.SignatureAlgorithm{jose.RS256})
	if err != nil {
		return "", 0, err
	}
	if len(parsed.Headers) == 0 {
		return "", 0, errors.New("missing token headers")
	}
	key, err := a.jwks.GetKey(ctx, parsed.Headers[0].KeyID)
	if err != nil {
		return "", 0, err
	}
	var std jwt.Claims
	custom := map[string]any{}
	if err := parsed.Claims(key, &std, &custom); err != nil {
		return "", 0, err
	}
	expected := jwt.Expected{Time: time.Now()}
	if a.cfg.OIDCIssuer != "" {
		expected.Issuer = a.cfg.OIDCIssuer
	}
	if a.cfg.OIDCAudience != "" {
		expected.AnyAudience = parseAudiences(a.cfg.OIDCAudience)
	}
	if err := std.ValidateWithLeeway(expected, 30*time.Second); err != nil {
		return "", 0, err
	}
	if std.Subject == "" {
		return "", 0, errors.New("missing subject claim")
	}
	at, ok := numericClaim(custom["auth_time"])
	if !ok {
		return "", 0, errors.New("missing auth_time claim")
	}
	return std.Subject, at, nil
}

// numericClaim coerces a JSON number claim (float64 from stdlib unmarshal, or
// json.Number / int64) to int64.
func numericClaim(v any) (int64, bool) {
	switch n := v.(type) {
	case float64:
		return int64(n), true
	case int64:
		return n, true
	case json.Number:
		if i, err := n.Int64(); err == nil {
			return i, true
		}
	}
	return 0, false
}

// parseAudiences splits a comma-separated OIDCAudience config into a jwt.Audience
// set. Multiple values let one backend accept tokens minted for several OAuth
// clients (the web app and the desktop client carry different `aud`); go-jose's
// AnyAudience passes when the token's audience matches any one entry. Whitespace
// around entries is trimmed and empties dropped, so a plain single value behaves
// exactly as before.
func parseAudiences(raw string) jwt.Audience {
	parts := strings.Split(raw, ",")
	out := make(jwt.Audience, 0, len(parts))
	for _, p := range parts {
		if s := strings.TrimSpace(p); s != "" {
			out = append(out, s)
		}
	}
	return out
}

func claimsFromJWT(std jwt.Claims, custom map[string]any) (*Claims, error) {
	if std.Subject == "" {
		return nil, errors.New("missing subject claim")
	}
	c := &Claims{UserID: std.Subject}
	if g, ok := custom["groups"].([]any); ok {
		for _, item := range g {
			if s, ok := item.(string); ok {
				c.Groups = append(c.Groups, s)
			}
		}
	}
	return c, nil
}

func bearerToken(r *http.Request) (string, bool) {
	h := r.Header.Get("Authorization")
	const prefix = "Bearer "
	if !strings.HasPrefix(h, prefix) {
		return "", false
	}
	tok := strings.TrimPrefix(h, prefix)
	if tok == "" {
		return "", false
	}
	return tok, true
}

func unauthorized(w http.ResponseWriter, reason string) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.Header().Set("WWW-Authenticate", `Bearer realm="cdrop"`)
	w.WriteHeader(http.StatusUnauthorized)
	_ = json.NewEncoder(w).Encode(map[string]string{
		"error":  "unauthorized",
		"reason": reason,
	})
}

// DeriveHS256Key turns a configured secret of any length into a fixed 32-byte
// HMAC key. HS256 requires >= 32 bytes; hashing guarantees that (and keeps the
// minting and verifying sides in lockstep) so a short CDROP_HS256_SECRET can't
// make shortcut-token signing fail. Both signing (httpapi) and verifying use it.
func DeriveHS256Key(secret string) []byte {
	sum := sha256.Sum256([]byte(secret))
	return sum[:]
}

// DeriveSessionTokenKey derives the HMAC key for cdrop's self-signed session
// access tokens from CDROP_SESSION_SECRET. Domain-separated from both the at-rest
// refresh-token AES key (plain sha256(secret), in httpapi) and the shortcut-token
// key (DeriveHS256Key) so one secret yields three independent keys — a session
// token can never validate as a shortcut token, or vice versa. Empty secret →
// nil, which disables both minting and verifying (dev / unconfigured prod).
func DeriveSessionTokenKey(secret string) []byte {
	if secret == "" {
		return nil
	}
	sum := sha256.Sum256([]byte("cdrop-session-jwt\x00" + secret))
	return sum[:]
}

// SanitizeDeviceName enforces the global ASCII-only device-name policy. Device
// names ride in the X-Device-Name HTTP header, which can't carry non-ASCII
// reliably (and browser fetch rejects such header values outright), so the name
// is restricted to printable ASCII everywhere. Here we keep only printable ASCII
// (0x20–0x7E), trim, and cap the length as a server-side backstop; clients also
// validate the name up front for a clear error. Empty after sanitising → the
// caller falls back to a UA-derived default.
func SanitizeDeviceName(raw string) string {
	var b strings.Builder
	for _, r := range raw {
		if r >= 0x20 && r <= 0x7E {
			b.WriteRune(r)
		}
	}
	name := strings.TrimSpace(b.String())
	if len(name) > 64 {
		name = strings.TrimSpace(name[:64])
	}
	return name
}

// normalizeDeviceType whitelists the client-declared X-Device-Type so a device
// row only ever carries a known kind; anything unrecognised (incl. empty) falls
// back to "browser", the default web client. Desktop clients send macos/windows;
// "ios" is reserved for a future native iOS client (the Shortcut never sends it).
func normalizeDeviceType(raw string) string {
	switch t := strings.ToLower(strings.TrimSpace(raw)); t {
	case "macos", "windows", "linux", "ios", "browser":
		return t
	default:
		return "browser"
	}
}