package authn

import (
	"context"
	"crypto/ed25519"
	"crypto/subtle"
	"fmt"
	"log"
	"time"

	"sourcecraft.dev/bigbes/sr-ht-core/auth"
	"github.com/dolthub/dolt/go/libraries/doltcore/creds"
	jose "gopkg.in/go-jose/go-jose.v2"
	"gopkg.in/go-jose/go-jose.v2/jwt"
)

// jwtSubjectPrefix is the fixed prefix dolt puts in the JWT "sub" claim,
// followed by the base32 key id: "doltClientCredentials/<kid>". See dolt's
// creds.DoltCreds.RPCCreds.
const jwtSubjectPrefix = "doltClientCredentials/"

// jwtLeeway is the clock-skew tolerance applied to the exp claim. dolt mints
// 30-second tokens; a minute of leeway absorbs modest clock drift between the
// client and this server without meaningfully extending a token's usefulness.
const jwtLeeway = 1 * time.Minute

// KeyStore looks up dolt Ed25519 public keys by their key id and records key
// usage. It is implemented by the db package (concurrently); authn defines it
// here to stay independent of the database. A key id is
// creds.PubKeyToKIDStr(pubkey): base32(SHA-512/224(pubkey)) over dolt's custom
// alphabet.
type KeyStore interface {
	// ByKID returns the 32-byte Ed25519 public key registered under kid and the
	// username of its owner. It returns an error if no such key exists.
	ByKID(ctx context.Context, kid string) (pubkey []byte, username string, err error)
	// TouchLastUsed records that the key identified by kid was just used.
	TouchLastUsed(ctx context.Context, kid string) error
}

// ResolveDoltJWT resolves the caller for a dolt keypair Bearer token: an EdDSA
// JWS minted by `dolt clone/push` from a `dolt creds`/`dolt login` keypair. It:
//
//   - parses the JWS and requires alg == EdDSA (defeats alg-confusion / "none");
//   - reads the kid header and looks the public key up via keys.ByKID;
//   - cross-checks kid == creds.PubKeyToKIDStr(pubkey) — the same derivation
//     dolt uses — so a corrupt (kid, pubkey) pairing cannot be trusted;
//   - verifies the EdDSA signature with the stored public key;
//   - requires aud == expectedAud and a valid exp (with jwtLeeway), and
//     sub == "doltClientCredentials/<kid>"; the iss claim is deliberately
//     IGNORED (dolt hardcodes iss = "dolt-client.dolthub.com");
//   - resolves the owning user via the meta backend and records key usage.
//
// Permanent rejections wrap ErrInvalidToken; a backend failure (user lookup)
// is returned unwrapped so callers treat it as transient. A TouchLastUsed
// failure is non-fatal (bookkeeping only) and is logged, not returned.
func ResolveDoltJWT(ctx context.Context, rawJWT, expectedAud string, keys KeyStore) (*auth.AuthContext, error) {
	tok, err := jwt.ParseSigned(rawJWT)
	if err != nil {
		return nil, fmt.Errorf("%w: parsing JWS: %v", ErrInvalidToken, err)
	}
	if len(tok.Headers) != 1 {
		return nil, fmt.Errorf("%w: expected exactly one signature, got %d", ErrInvalidToken, len(tok.Headers))
	}
	hdr := tok.Headers[0]
	if hdr.Algorithm != string(jose.EdDSA) {
		return nil, fmt.Errorf("%w: unexpected signature algorithm %q (want EdDSA)", ErrInvalidToken, hdr.Algorithm)
	}
	kid := hdr.KeyID
	if kid == "" {
		return nil, fmt.Errorf("%w: missing kid header", ErrInvalidToken)
	}

	pubkey, username, err := keys.ByKID(ctx, kid)
	if err != nil {
		return nil, fmt.Errorf("%w: unknown key id %q: %v", ErrInvalidToken, kid, err)
	}
	if len(pubkey) != ed25519.PublicKeySize {
		return nil, fmt.Errorf("%w: stored public key for %q has wrong size %d", ErrInvalidToken, kid, len(pubkey))
	}
	// Integrity: the key id must be the canonical derivation of this public
	// key, computed exactly as dolt does. Guards against a mismatched DB row.
	if subtle.ConstantTimeCompare([]byte(creds.PubKeyToKIDStr(pubkey)), []byte(kid)) != 1 {
		return nil, fmt.Errorf("%w: key id %q does not match its stored public key", ErrInvalidToken, kid)
	}

	var claims jwt.Claims
	if err := tok.Claims(ed25519.PublicKey(pubkey), &claims); err != nil {
		return nil, fmt.Errorf("%w: signature verification failed: %v", ErrInvalidToken, err)
	}

	// Validate aud + exp/nbf (with leeway). Issuer and Subject are intentionally
	// left unset in Expected: iss is ignored, and sub is checked separately
	// below against the exact "doltClientCredentials/<kid>" form.
	if err := claims.ValidateWithLeeway(jwt.Expected{
		Audience: jwt.Audience{expectedAud},
		Time:     nowFn(),
	}, jwtLeeway); err != nil {
		return nil, fmt.Errorf("%w: claim validation failed: %v", ErrInvalidToken, err)
	}
	if claims.Subject != jwtSubjectPrefix+kid {
		return nil, fmt.Errorf("%w: subject %q is not %s%s", ErrInvalidToken, claims.Subject, jwtSubjectPrefix, kid)
	}

	var ac auth.AuthContext
	if err := meta.LookupUser(ctx, username, &ac); err != nil {
		return nil, fmt.Errorf("looking up user %q: %w", username, err)
	}
	ac.AuthMethod = AuthMethodDoltKey

	if err := keys.TouchLastUsed(ctx, kid); err != nil {
		// Non-fatal: the caller is already authenticated; last_used is display
		// metadata. Log and continue rather than failing the clone/push.
		log.Printf("authn: TouchLastUsed(%q): %v", kid, err)
	}
	return &ac, nil
}