package oidc

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"encoding/json"
	"net/http"
	"strings"
	"time"

	"keycape/internal/adapters/tenantengine"
	"keycape/internal/domain"
	profileerrors "keycape/internal/errors"
	"keycape/internal/server/telemetry"
)

// TokenHandler implements POST /token.
type TokenHandler struct {
	ClientConfig  map[string]*domain.Client
	Sessions      *SessionStore
	Users         domain.UserRepository
	SigningKey    *rsa.PrivateKey
	Issuer        string
	TokenLifetime time.Duration
	Emitter       telemetry.Emitter
	// TenantEngine sources the optional tenant_roles claim (KEY-WP-0005-T02).
	// Nil disables it entirely -- token issuance never depends on it.
	TenantEngine *tenantengine.Client
}

// tokenResponse is the JSON body returned on a successful token exchange.
type tokenResponse struct {
	AccessToken string `json:"access_token"`
	TokenType   string `json:"token_type"`
	ExpiresIn   int    `json:"expires_in"`
	IDToken     string `json:"id_token,omitempty"`
}

// ServeHTTP handles POST /token.
func (h *TokenHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	if err := r.ParseForm(); err != nil {
		http.Error(w, "invalid form body", http.StatusBadRequest)
		return
	}

	grantType := r.FormValue("grant_type")
	if grantType == "client_credentials" {
		h.serveClientCredentials(w, r)
		return
	}
	clientID := r.FormValue("client_id")
	code := r.FormValue("code")
	codeVerifier := r.FormValue("code_verifier")

	// 1. Validate grant_type.
	if grantType != "authorization_code" {
		profileerrors.FeatureNotSupported(
			"only grant_type=authorization_code is supported",
			"grant_type="+grantType,
		).Write(w, http.StatusBadRequest)
		return
	}

	// 2. Validate client exists (basic check; secret auth delegated to future work).
	if _, ok := h.ClientConfig[clientID]; !ok {
		profileerrors.InvalidProfileUsage("unknown client_id", "client_id").
			Write(w, http.StatusBadRequest)
		return
	}

	// 3. Look up PKCE session.
	sess, ok := h.Sessions.Get(code)
	if !ok {
		profileerrors.InvalidProfileUsage(
			"authorization code not found or expired",
			"code",
		).Write(w, http.StatusBadRequest)
		return
	}

	// Verify client_id matches the session.
	if sess.ClientID != clientID {
		profileerrors.InvalidProfileUsage(
			"client_id does not match the authorization code",
			"client_id",
		).Write(w, http.StatusBadRequest)
		return
	}

	// 4. Verify PKCE code_verifier.
	if !verifyPKCE(codeVerifier, sess.PKCEChallenge) {
		profileerrors.InvalidProfileUsage(
			"code_verifier does not match code_challenge",
			"code_verifier",
		).Write(w, http.StatusBadRequest)
		return
	}

	// 5. Look up user.
	user, err := h.Users.LookupUser(ctx, sess.Username)
	if err != nil {
		http.Error(w, "user not found", http.StatusInternalServerError)
		return
	}
	if isSuspended(user) {
		h.Sessions.Delete(code)
		profileerrors.RejectedForSafety(
			"account is suspended",
			"account_lifecycle",
		).Write(w, http.StatusForbidden)
		return
	}

	// 6. Build JWT claims.
	now := time.Now()
	exp := now.Add(h.TokenLifetime)

	claims := map[string]interface{}{
		"iss": h.Issuer,
		"sub": user.ID,
		"aud": clientID,
		"exp": exp.Unix(),
		"iat": now.Unix(),
	}
	if sess.Nonce != "" {
		claims["nonce"] = sess.Nonce
	}

	scopeSet := make(map[string]bool)
	for _, s := range sess.Scopes {
		scopeSet[s] = true
	}

	if scopeSet["profile"] {
		claims["preferred_username"] = user.Username
	}
	if scopeSet["email"] {
		claims["email"] = user.Email
	}

	// Core claims required by net-kingdom/canon/standards/iam-profile_v0.3.md
	// for every production token -- not scope-gated, unlike the recommended
	// human claims above (KEY-WP-0005-T01).
	tenant := effectiveTenant(user)
	claims["tenant"] = tenant
	claims["principal_type"] = "human"
	claims["groups"] = nonNilStrings(user.Groups)
	claims["roles"] = nonNilStrings(user.Roles)
	claims["assurance"] = assuranceClaim(sess.MFAVerified, now)

	// Optional cached tenant_roles claim (KEY-WP-0005-T02). Fails open --
	// see internal/adapters/tenantengine's package doc for why this is the
	// one place in the whole tenant_roles design where that's correct.
	if roles, ok := h.TenantEngine.Roles(ctx, tenant); ok {
		claims["tenant_roles"] = roles
	}

	// 7. Sign JWT with RSA-SHA256.
	kid := "key-1" // static kid for v0.1
	jwtToken, err := buildJWT(claims, kid, h.SigningKey)
	if err != nil {
		http.Error(w, "failed to build JWT", http.StatusInternalServerError)
		return
	}

	// 8. Delete used PKCE session (prevent replay).
	h.Sessions.Delete(code)

	// 9. Build response.
	resp := tokenResponse{
		AccessToken: jwtToken,
		TokenType:   "Bearer",
		ExpiresIn:   int(h.TokenLifetime.Seconds()),
		IDToken:     jwtToken,
	}

	// 10. Emit token_issued telemetry.
	h.Emitter.Emit(ctx, telemetry.Event{
		Timestamp: time.Now(),
		EventType: telemetry.EventTokenIssued,
		ClientID:  clientID,
		Endpoint:  "/token",
		Result:    "success",
		Scopes:    sess.Scopes,
		GrantType: grantType,
	})

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	_ = json.NewEncoder(w).Encode(resp)
}

func (h *TokenHandler) serveClientCredentials(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	clientID, clientSecret, ok := r.BasicAuth()
	if !ok {
		profileerrors.InvalidProfileUsage("client_secret_basic authentication required", "Authorization").
			Write(w, http.StatusUnauthorized)
		return
	}
	client, ok := h.ClientConfig[clientID]
	if !ok || client.ClientType != "confidential" || !containsString(client.GrantTypes, "client_credentials") {
		profileerrors.InvalidProfileUsage("invalid confidential client", "client_id").
			Write(w, http.StatusUnauthorized)
		return
	}
	presentedDigest := sha256.Sum256([]byte(clientSecret))
	expectedDigest := sha256.Sum256([]byte(client.ClientSecret))
	if client.ClientSecret == "" ||
		subtle.ConstantTimeCompare(presentedDigest[:], expectedDigest[:]) != 1 {
		profileerrors.InvalidProfileUsage("invalid client authentication", "Authorization").
			Write(w, http.StatusUnauthorized)
		return
	}

	requestedScopes := strings.Fields(r.FormValue("scope"))
	if len(requestedScopes) == 0 {
		requestedScopes = append([]string(nil), client.AllowedScopes...)
	}
	for _, scope := range requestedScopes {
		if !containsString(client.AllowedScopes, scope) {
			profileerrors.InvalidProfileUsage("requested scope is not allowed", "scope").
				Write(w, http.StatusBadRequest)
			return
		}
	}

	now := time.Now()
	tokenLifetime := h.TokenLifetime
	if client.TokenLifetime > 0 {
		tokenLifetime = client.TokenLifetime
	}
	claims := map[string]interface{}{
		"iss":            h.Issuer,
		"sub":            client.ServiceSubject,
		"aud":            clientID,
		"exp":            now.Add(tokenLifetime).Unix(),
		"iat":            now.Unix(),
		"tenant":         client.Tenant,
		"principal_type": "service",
		"groups":         []string{},
		"roles":          nonNilStrings(client.Roles),
		"scope":          strings.Join(requestedScopes, " "),
		"assurance": map[string]interface{}{
			"level": "aal1", "methods": []string{"client_secret"},
			"mfa": false, "source": "key-cape", "at": now.Unix(),
		},
	}
	if roles, ok := h.TenantEngine.Roles(ctx, client.Tenant); ok {
		claims["tenant_roles"] = roles
	}
	jwtToken, err := buildJWT(claims, "key-1", h.SigningKey)
	if err != nil {
		http.Error(w, "failed to build JWT", http.StatusInternalServerError)
		return
	}
	h.Emitter.Emit(ctx, telemetry.Event{
		Timestamp: now, EventType: telemetry.EventTokenIssued, ClientID: clientID,
		Endpoint: "/token", Result: "success", Scopes: requestedScopes,
		GrantType: "client_credentials",
	})
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	_ = json.NewEncoder(w).Encode(tokenResponse{
		AccessToken: jwtToken, TokenType: "Bearer",
		ExpiresIn: int(tokenLifetime.Seconds()),
	})
}

func containsString(values []string, wanted string) bool {
	for _, value := range values {
		if value == wanted {
			return true
		}
	}
	return false
}

func isSuspended(user *domain.User) bool {
	return containsString(user.Groups, "netkingdom-suspended")
}

// ---------------------------------------------------------------------------
// IAM Profile core claims (KEY-WP-0005-T01)
// ---------------------------------------------------------------------------

// defaultTenant is the fallback tenant claim for users with no explicit
// Tenant assignment yet. This workstation currently operates a single
// tenant (tenant:coulomb, ADR-0006); later tenants (e.g. tenant:friendly:binky,
// ADR-0013) require an explicit domain.User.Tenant value -- this default
// never silently assigns a user to a tenant other than the platform's
// original one.
const defaultTenant = "tenant:coulomb"

// effectiveTenant resolves the tenant claim for a user, falling back to
// defaultTenant when the user has no explicit tenant assignment. The IAM
// Profile requires a non-empty tenant claim on every token.
func effectiveTenant(user *domain.User) string {
	if user.Tenant != "" {
		return user.Tenant
	}
	return defaultTenant
}

// nonNilStrings returns s, or an empty (non-nil) slice if s is nil, so the
// claim always serializes as `[]`, never `null` -- the profile requires
// groups/roles to be present, "possibly empty", not absent.
func nonNilStrings(s []string) []string {
	if s == nil {
		return []string{}
	}
	return s
}

// assuranceClaim builds the profile's `assurance` object from whether MFA
// was actually verified during this authorization (session.MFAVerified),
// not from static enrollment state -- a user who has MFA enrolled but
// wasn't challenged for it in this particular flow gets aal1, not aal2.
func assuranceClaim(mfaVerified bool, at time.Time) map[string]interface{} {
	level := "aal1"
	methods := []string{"pwd"}
	if mfaVerified {
		level = "aal2"
		methods = append(methods, "otp")
	}
	return map[string]interface{}{
		"level":   level,
		"methods": methods,
		"mfa":     mfaVerified,
		"source":  "key-cape",
		"at":      at.Unix(),
	}
}

// ---------------------------------------------------------------------------
// PKCE verification
// ---------------------------------------------------------------------------

// verifyPKCE checks BASE64URL(SHA256(verifier)) == challenge (S256 method).
func verifyPKCE(verifier, challenge string) bool {
	h := sha256.New()
	h.Write([]byte(verifier))
	computed := base64.RawURLEncoding.EncodeToString(h.Sum(nil))
	return computed == challenge
}

// ---------------------------------------------------------------------------
// JWT construction (stdlib only — no external JWT library)
// ---------------------------------------------------------------------------

type jwtHeader struct {
	Alg string `json:"alg"`
	Typ string `json:"typ"`
	Kid string `json:"kid"`
}

// buildJWT constructs and signs a JWT using RSA-SHA256 with the standard library.
// Format: base64url(header) + "." + base64url(payload) + "." + base64url(signature)
func buildJWT(claims map[string]interface{}, kid string, key *rsa.PrivateKey) (string, error) {
	// Header.
	hdr := jwtHeader{Alg: "RS256", Typ: "JWT", Kid: kid}
	hdrJSON, err := json.Marshal(hdr)
	if err != nil {
		return "", err
	}
	hdrB64 := base64.RawURLEncoding.EncodeToString(hdrJSON)

	// Payload.
	payloadJSON, err := json.Marshal(claims)
	if err != nil {
		return "", err
	}
	payloadB64 := base64.RawURLEncoding.EncodeToString(payloadJSON)

	// Signing input.
	signingInput := hdrB64 + "." + payloadB64

	// Digest.
	digest := sha256.Sum256([]byte(signingInput))

	// Sign with PKCS1v15 / SHA256.
	sig, err := rsa.SignPKCS1v15(rand.Reader, key, crypto.SHA256, digest[:])
	if err != nil {
		return "", err
	}
	sigB64 := base64.RawURLEncoding.EncodeToString(sig)

	return strings.Join([]string{hdrB64, payloadB64, sigB64}, "."), nil
}