package auth

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math/big"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

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

const maxOIDCResponseBytes = 1 << 20

type OIDCConfig struct {
	Issuer                    string
	Audience                  string
	TenantID                  string
	RolePrefix                string
	RequiredScopes            []string
	JWKSCacheTTL              time.Duration
	IntrospectionEnabled      bool
	IntrospectionClientID     string
	IntrospectionClientSecret string
	HTTPClient                *http.Client
}

type OIDCVerifier struct {
	config                OIDCConfig
	client                *http.Client
	mu                    sync.Mutex
	keys                  map[string]any
	expiresAt             time.Time
	introspectionEndpoint string
}

type discoveryDocument struct {
	Issuer                string `json:"issuer"`
	JWKSURI               string `json:"jwks_uri"`
	IntrospectionEndpoint string `json:"introspection_endpoint"`
}

type jsonWebKeySet struct {
	Keys []jsonWebKey `json:"keys"`
}

type jsonWebKey struct {
	KID string `json:"kid"`
	KTY string `json:"kty"`
	Use string `json:"use"`
	Alg string `json:"alg"`
	N   string `json:"n"`
	E   string `json:"e"`
	Crv string `json:"crv"`
	X   string `json:"x"`
	Y   string `json:"y"`
}

func NewOIDCVerifier(config OIDCConfig) (*OIDCVerifier, error) {
	config.Issuer = strings.TrimRight(strings.TrimSpace(config.Issuer), "/")
	config.Audience = strings.TrimSpace(config.Audience)
	config.TenantID = strings.TrimSpace(config.TenantID)
	config.RolePrefix = strings.TrimSpace(config.RolePrefix)
	if err := validatePublicURL(config.Issuer); err != nil || config.Audience == "" || config.TenantID == "" || config.RolePrefix == "" {
		return nil, errors.New("issuer, audience, tenant and role prefix are required")
	}
	if config.IntrospectionEnabled && (strings.TrimSpace(config.IntrospectionClientID) == "" || config.IntrospectionClientSecret == "") {
		return nil, errors.New("introspection client credentials are required when introspection is enabled")
	}
	if config.JWKSCacheTTL <= 0 {
		config.JWKSCacheTTL = 5 * time.Minute
	}
	client := config.HTTPClient
	if client == nil {
		client = &http.Client{Timeout: 10 * time.Second, CheckRedirect: func(_ *http.Request, _ []*http.Request) error {
			return http.ErrUseLastResponse
		}}
	}
	return &OIDCVerifier{config: config, client: client, keys: map[string]any{}}, nil
}

func (v *OIDCVerifier) Verify(ctx context.Context, raw string) (*User, error) {
	parser := jwt.NewParser(jwt.WithValidMethods([]string{"RS256", "ES256"}))
	unverified, _, err := parser.ParseUnverified(raw, jwt.MapClaims{})
	if err != nil || unverified == nil {
		return nil, oidcUnauthorized("token envelope is invalid", err)
	}
	kid, _ := unverified.Header["kid"].(string)
	if kid == "" {
		return nil, oidcUnauthorized("kid is missing", nil)
	}
	key, err := v.key(ctx, kid)
	if err != nil {
		return nil, oidcUnauthorized("signing key lookup failed", err)
	}
	token, err := jwt.Parse(raw, func(token *jwt.Token) (any, error) {
		if token.Header["kid"] != kid {
			return nil, ErrUnauthorized
		}
		return key, nil
	}, jwt.WithValidMethods([]string{"RS256", "ES256"}), jwt.WithIssuer(v.config.Issuer),
		jwt.WithAudience(v.config.Audience), jwt.WithExpirationRequired(), jwt.WithLeeway(30*time.Second))
	if err != nil || !token.Valid {
		return nil, oidcUnauthorized("signature or registered claims are invalid", err)
	}
	claims, ok := token.Claims.(jwt.MapClaims)
	if !ok || stringClaim(claims, "sub") == "" || stringClaim(claims, "tid") != v.config.TenantID {
		return nil, oidcUnauthorized("stable identity claims are invalid", nil)
	}
	if _, ok := numericDateClaim(claims["nbf"]); !ok {
		return nil, oidcUnauthorized("nbf is missing", nil)
	}
	scopes := scopeClaims(claims)
	if !containsAll(scopes, v.config.RequiredScopes) {
		return nil, oidcUnauthorized("required scope is missing", nil)
	}
	roles := gatewayRoles(stringSliceClaim(claims, "roles"), v.config.RolePrefix)
	if len(roles) == 0 {
		roles = gatewayRoles(stringSliceClaim(claims, "role"), v.config.RolePrefix)
	}
	if len(roles) == 0 {
		return nil, oidcUnauthorized("mapped role is missing", nil)
	}
	if v.config.IntrospectionEnabled {
		active, err := v.introspect(ctx, raw)
		if err != nil || !active {
			return nil, oidcUnauthorized("token is inactive", err)
		}
	}
	username := stringClaim(claims, "preferred_username")
	if username == "" {
		username = stringClaim(claims, "username")
	}
	return &User{
		ID: stringClaim(claims, "sub"), Username: username, Roles: roles,
		TenantID: v.config.TenantID, Source: "oidc",
	}, nil
}

func oidcUnauthorized(reason string, cause error) error {
	if cause != nil {
		return fmt.Errorf("%w: %s: %v", ErrUnauthorized, reason, cause)
	}
	return fmt.Errorf("%w: %s", ErrUnauthorized, reason)
}

func (v *OIDCVerifier) key(ctx context.Context, kid string) (any, error) {
	if err := v.refresh(ctx, false); err != nil {
		return nil, err
	}
	v.mu.Lock()
	key, ok := v.keys[kid]
	v.mu.Unlock()
	if ok {
		return key, nil
	}
	if err := v.refresh(ctx, true); err != nil {
		return nil, err
	}
	v.mu.Lock()
	defer v.mu.Unlock()
	key, ok = v.keys[kid]
	if !ok {
		return nil, errors.New("signing key not found")
	}
	return key, nil
}

func (v *OIDCVerifier) refresh(ctx context.Context, force bool) error {
	v.mu.Lock()
	defer v.mu.Unlock()
	if !force && len(v.keys) > 0 && time.Now().Before(v.expiresAt) {
		return nil
	}
	discoveryURL := v.config.Issuer + "/.well-known/openid-configuration"
	var discovery discoveryDocument
	if err := v.fetchJSON(ctx, discoveryURL, &discovery); err != nil {
		return err
	}
	if discovery.Issuer != v.config.Issuer || validatePublicURL(discovery.JWKSURI) != nil {
		return errors.New("OIDC discovery metadata is invalid")
	}
	if v.config.IntrospectionEnabled && validatePublicURL(discovery.IntrospectionEndpoint) != nil {
		return errors.New("OIDC introspection metadata is invalid")
	}
	var set jsonWebKeySet
	if err := v.fetchJSON(ctx, discovery.JWKSURI, &set); err != nil {
		return err
	}
	keys := make(map[string]any, len(set.Keys))
	for _, item := range set.Keys {
		if item.KID == "" || item.Use != "" && item.Use != "sig" {
			continue
		}
		key, err := item.publicKey()
		if err == nil {
			keys[item.KID] = key
		}
	}
	if len(keys) == 0 {
		return errors.New("OIDC JWKS contains no supported signing keys")
	}
	v.keys = keys
	v.expiresAt = time.Now().Add(v.config.JWKSCacheTTL)
	v.introspectionEndpoint = discovery.IntrospectionEndpoint
	return nil
}

func (v *OIDCVerifier) introspect(ctx context.Context, raw string) (bool, error) {
	v.mu.Lock()
	endpoint := v.introspectionEndpoint
	v.mu.Unlock()
	if endpoint == "" {
		return false, errors.New("OIDC introspection endpoint is unavailable")
	}
	form := url.Values{"token": {raw}, "token_type_hint": {"access_token"}}
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, strings.NewReader(form.Encode()))
	if err != nil {
		return false, err
	}
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	request.Header.Set("Accept", "application/json")
	request.SetBasicAuth(v.config.IntrospectionClientID, v.config.IntrospectionClientSecret)
	response, err := v.client.Do(request)
	if err != nil {
		return false, err
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, maxOIDCResponseBytes))
		return false, fmt.Errorf("OIDC introspection returned HTTP %d", response.StatusCode)
	}
	var result struct {
		Active bool `json:"active"`
	}
	decoder := json.NewDecoder(io.LimitReader(response.Body, maxOIDCResponseBytes))
	if err := decoder.Decode(&result); err != nil {
		return false, errors.New("OIDC introspection response is invalid")
	}
	return result.Active, nil
}

func (v *OIDCVerifier) fetchJSON(ctx context.Context, endpoint string, output any) error {
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
	if err != nil {
		return err
	}
	request.Header.Set("Accept", "application/json")
	response, err := v.client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		return fmt.Errorf("OIDC endpoint returned HTTP %d", response.StatusCode)
	}
	payload, err := io.ReadAll(io.LimitReader(response.Body, maxOIDCResponseBytes+1))
	if err != nil || len(payload) > maxOIDCResponseBytes {
		return errors.New("OIDC response is invalid")
	}
	decoder := json.NewDecoder(strings.NewReader(string(payload)))
	return decoder.Decode(output)
}

func (j jsonWebKey) publicKey() (any, error) {
	switch {
	case j.KTY == "RSA" && (j.Alg == "" || j.Alg == "RS256"):
		n, err := decodeBigInt(j.N)
		if err != nil {
			return nil, err
		}
		e, err := decodeBigInt(j.E)
		if err != nil || !e.IsInt64() || e.Int64() <= 0 {
			return nil, errors.New("invalid RSA exponent")
		}
		return &rsa.PublicKey{N: n, E: int(e.Int64())}, nil
	case j.KTY == "EC" && j.Crv == "P-256" && (j.Alg == "" || j.Alg == "ES256"):
		x, err := decodeBigInt(j.X)
		if err != nil {
			return nil, err
		}
		y, err := decodeBigInt(j.Y)
		if err != nil || !elliptic.P256().IsOnCurve(x, y) {
			return nil, errors.New("invalid EC point")
		}
		return &ecdsa.PublicKey{Curve: elliptic.P256(), X: x, Y: y}, nil
	default:
		return nil, errors.New("unsupported JWK")
	}
}

func decodeBigInt(value string) (*big.Int, error) {
	payload, err := base64.RawURLEncoding.DecodeString(value)
	if err != nil || len(payload) == 0 {
		return nil, errors.New("invalid JWK integer")
	}
	return new(big.Int).SetBytes(payload), nil
}

func validatePublicURL(raw string) error {
	parsed, err := url.Parse(raw)
	if err != nil || parsed.Host == "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" {
		return errors.New("invalid URL")
	}
	if parsed.Scheme == "https" {
		return nil
	}
	if parsed.Scheme == "http" && (parsed.Hostname() == "127.0.0.1" || parsed.Hostname() == "localhost") {
		return nil
	}
	return errors.New("URL must use HTTPS")
}

func numericDateClaim(value any) (time.Time, bool) {
	switch typed := value.(type) {
	case float64:
		return time.Unix(int64(typed), 0), true
	case json.Number:
		seconds, err := typed.Int64()
		return time.Unix(seconds, 0), err == nil
	default:
		return time.Time{}, false
	}
}

func scopeClaims(claims jwt.MapClaims) []string {
	result := make([]string, 0)
	seen := map[string]bool{}
	for _, scope := range strings.Fields(stringClaim(claims, "scope")) {
		if !seen[scope] {
			result = append(result, scope)
			seen[scope] = true
		}
	}
	for _, scope := range stringSliceClaim(claims, "scp") {
		if !seen[scope] {
			result = append(result, scope)
			seen[scope] = true
		}
	}
	return result
}

func containsAll(actual, required []string) bool {
	set := make(map[string]struct{}, len(actual))
	for _, item := range actual {
		set[item] = struct{}{}
	}
	for _, item := range required {
		if _, ok := set[item]; !ok {
			return false
		}
	}
	return true
}

func gatewayRoles(external []string, prefix string) []string {
	allowed := map[string]bool{"user": true, "creator": true, "operator": true, "manager": true, "admin": true}
	roles := make([]string, 0, len(external))
	seen := map[string]bool{}
	for _, role := range external {
		if !strings.HasPrefix(role, prefix) {
			continue
		}
		local := strings.TrimPrefix(role, prefix)
		if allowed[local] && !seen[local] {
			roles = append(roles, local)
			seen[local] = true
		}
	}
	return roles
}