// Package auth verifies Keycloak-issued JWTs.
//
// JWKS discovery is lazy: we don't hit `.well-known/openid-configuration`
// at startup because pod-to-public-internet TLS may be slow during pod
// boot and we don't want a healthy DB-only deployment to crash because
// auth.gosec.cloud hiccupped. Discovery happens on first token validation
// and the result is cached. We retry on every fresh request until it
// succeeds.
package auth

import (
	"context"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/coreos/go-oidc/v3/oidc"
	"github.com/gofiber/fiber/v2"
)

type UserClaims struct {
	Subject     string   `json:"sub"`
	Email       string   `json:"email"`
	DisplayName string   `json:"name"`
	GivenName   string   `json:"given_name"`
	FamilyName  string   `json:"family_name"`
	TenantID    string   `json:"tenant_id"`
	Roles       []string // realm roles + flattened client roles
}

type Verifier struct {
	issuer        string // the canonical issuer claim Keycloak emits
	discoveryURL  string // optional: point discovery at an internal URL
	allowedAudCSV string

	mu sync.Mutex
	v  *oidc.IDTokenVerifier // nil until first successful discovery
}

// NewVerifier returns a Verifier that validates tokens whose `iss` matches
// `issuer`. If `discoveryURL` is non-empty, OIDC discovery is performed
// against that URL instead — useful when the canonical issuer is a public
// hostname (auth.gosec.cloud) but the pod can only reach Keycloak through
// an in-cluster service.
func NewVerifier(_ context.Context, issuer, discoveryURL, allowedAudiencesCSV string) (*Verifier, error) {
	if issuer == "" {
		return nil, errors.New("issuer is required")
	}
	return &Verifier{
		issuer:        issuer,
		discoveryURL:  discoveryURL,
		allowedAudCSV: allowedAudiencesCSV,
	}, nil
}

// ensure resolves the underlying OIDC verifier on demand. Cached after first success.
func (vr *Verifier) ensure(ctx context.Context) (*oidc.IDTokenVerifier, error) {
	vr.mu.Lock()
	defer vr.mu.Unlock()
	if vr.v != nil {
		return vr.v, nil
	}
	disc, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()

	target := vr.issuer
	if vr.discoveryURL != "" {
		// Tell go-oidc that "the issuer claim Keycloak puts in tokens" and
		// "the URL we hit for discovery" can differ — otherwise it errors
		// because the discovery doc's issuer field won't match what we
		// passed to NewProvider.
		disc = oidc.InsecureIssuerURLContext(disc, vr.issuer)
		target = vr.discoveryURL

		// In-cluster Keycloak presents a cert for `auth.gosec.cloud`, not
		// for `keycloak.keycloak.svc.cluster.local`. Ambient mesh policy
		// gates who can reach Keycloak; we don't need cert hostname
		// verification on top of that for the internal discovery hop.
		client := &http.Client{
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, //nolint:gosec
			},
			Timeout: 10 * time.Second,
		}
		disc = oidc.ClientContext(disc, client)
	}

	provider, err := oidc.NewProvider(disc, target)
	if err != nil {
		return nil, fmt.Errorf("oidc discovery: %w", err)
	}
	vr.v = provider.Verifier(&oidc.Config{
		// Custom audience check below; many Keycloak access tokens have
		// aud="account" which would fail strict OIDC checks.
		SkipClientIDCheck: true,
	})
	return vr.v, nil
}

func (vr *Verifier) Verify(ctx context.Context, raw string) (*UserClaims, error) {
	v, err := vr.ensure(ctx)
	if err != nil {
		return nil, err
	}
	tok, err := v.Verify(ctx, raw)
	if err != nil {
		return nil, err
	}

	var raw1 struct {
		Subject     string `json:"sub"`
		Email       string `json:"email"`
		Name        string `json:"name"`
		GivenName   string `json:"given_name"`
		FamilyName  string `json:"family_name"`
		TenantID    string `json:"tenant_id"`
		AzpAudience string `json:"azp"`
		RealmAccess struct {
			Roles []string `json:"roles"`
		} `json:"realm_access"`
		ResourceAccess map[string]struct {
			Roles []string `json:"roles"`
		} `json:"resource_access"`
	}
	if err := tok.Claims(&raw1); err != nil {
		return nil, fmt.Errorf("decode claims: %w", err)
	}

	if vr.allowedAudCSV != "" {
		if !contains(strings.Split(vr.allowedAudCSV, ","), raw1.AzpAudience) {
			return nil, errors.New("token azp not in allowed audiences")
		}
	}

	roles := append([]string{}, raw1.RealmAccess.Roles...)
	for _, ra := range raw1.ResourceAccess {
		roles = append(roles, ra.Roles...)
	}

	return &UserClaims{
		Subject:     raw1.Subject,
		Email:       raw1.Email,
		DisplayName: raw1.Name,
		GivenName:   raw1.GivenName,
		FamilyName:  raw1.FamilyName,
		TenantID:    raw1.TenantID,
		Roles:       roles,
	}, nil
}

func (vr *Verifier) Middleware() fiber.Handler {
	return func(c *fiber.Ctx) error {
		header := c.Get("Authorization")
		if !strings.HasPrefix(header, "Bearer ") {
			return c.Status(fiber.StatusUnauthorized).
				JSON(fiber.Map{"error": "missing bearer token"})
		}
		raw := strings.TrimPrefix(header, "Bearer ")
		claims, err := vr.Verify(c.UserContext(), raw)
		if err != nil {
			// Temporary diagnostic — log the actual verify error and the
			// token's header (kid, alg, typ) so we can tell signature vs
			// expiration vs key-not-found apart. Payload claims (sub, iss,
			// exp, aud) too so we can correlate with realm config.
			log.Printf("[auth] verify failed: %v | header=%s payload=%s",
				err, decodeJWTPart(raw, 0), decodeJWTPart(raw, 1))
			return c.Status(fiber.StatusUnauthorized).
				JSON(fiber.Map{"error": fmt.Sprintf("invalid token: %v", err)})
		}
		c.Locals("user", claims)
		return c.Next()
	}
}

// decodeJWTPart base64-decodes header (idx 0) or payload (idx 1) of a JWT
// and pretty-prints selected fields. Best-effort — returns the raw segment
// on parse failure. Never returns the signature segment.
func decodeJWTPart(token string, idx int) string {
	parts := strings.Split(token, ".")
	if idx >= len(parts) || idx > 1 {
		return "(missing)"
	}
	raw, err := base64.RawURLEncoding.DecodeString(parts[idx])
	if err != nil {
		return parts[idx][:min(len(parts[idx]), 40)] + "(b64err)"
	}
	var m map[string]any
	if err := json.Unmarshal(raw, &m); err != nil {
		return string(raw)[:min(len(raw), 80)]
	}
	keep := map[string]any{}
	for _, k := range []string{"kid", "alg", "typ", "iss", "aud", "azp", "sub", "exp", "iat"} {
		if v, ok := m[k]; ok {
			keep[k] = v
		}
	}
	b, _ := json.Marshal(keep)
	return string(b)
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func User(c *fiber.Ctx) *UserClaims {
	if u, ok := c.Locals("user").(*UserClaims); ok {
		return u
	}
	return nil
}

func contains(haystack []string, needle string) bool {
	for _, h := range haystack {
		if strings.TrimSpace(h) == needle {
			return true
		}
	}
	return false
}