// Package trusted extracts identity from JWTs that a protected listener
// has already verified and carries that identity through request contexts.
package trusted

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
)

type contextKey struct{}

// Identity is the claims projection consumed from a JWT that the protected
// listener has already verified. It is request identity context, not a JWT
// verification result, a generic AgentTeam user/member record, or authorization
// policy; callers must resolve those control-plane facts in their own owners.
type Identity struct {
	Subject           string   `json:"sub,omitempty"`
	Email             string   `json:"email,omitempty"`
	Name              string   `json:"name,omitempty"`
	PreferredUsername string   `json:"preferred_username,omitempty"`
	Groups            []string `json:"groups,omitempty"`
	Tenant            string   `json:"tenant,omitempty"`
	RawClaims         map[string]any
}

// TokenMapper maps an already trusted upstream identity and token to the bearer
// token required by a downstream in-process service. It does not verify JWT
// signatures or decide whether the identity is authorized for a resource.
type TokenMapper func(context.Context, Identity, string) (string, error)

// Options configures trusted-listener identity extraction middleware. Header
// must name a header protected from untrusted clients; AllowMock is only for
// explicitly configured development/test paths, and TokenMapper is an optional
// downstream credential bridge rather than an authenticator.
type Options struct {
	Enabled     bool
	Header      string
	AllowMock   bool
	TokenMapper TokenMapper
}

func Middleware(options Options) func(http.Handler) http.Handler {
	if !options.Enabled {
		return func(next http.Handler) http.Handler { return next }
	}
	if strings.TrimSpace(options.Header) == "" {
		options.Header = "Authorization"
	}
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			identity, token, err := IdentityFromRequest(r, options)
			if err != nil {
				http.Error(w, err.Error(), http.StatusUnauthorized)
				return
			}
			authToken := token
			if options.TokenMapper != nil {
				mapped, err := options.TokenMapper(r.Context(), identity, token)
				if err != nil {
					http.Error(w, err.Error(), http.StatusUnauthorized)
					return
				}
				authToken = mapped
			}
			clone := r.Clone(WithIdentity(r.Context(), identity))
			if authToken != "" {
				clone.Header.Set("Authorization", "Bearer "+authToken)
			}
			next.ServeHTTP(w, clone)
		})
	}
}

func WithIdentity(ctx context.Context, identity Identity) context.Context {
	return context.WithValue(ctx, contextKey{}, identity)
}

func IdentityFromContext(ctx context.Context) (Identity, bool) {
	identity, ok := ctx.Value(contextKey{}).(Identity)
	return identity, ok
}

// IdentityFromRequest extracts and decodes the AgentTeam-verified JWT from a request.
func IdentityFromRequest(r *http.Request, options Options) (Identity, string, error) {
	header := strings.TrimSpace(options.Header)
	if header == "" {
		header = "Authorization"
	}
	token := bearerToken(r.Header.Get(header))
	if token == "" && !strings.EqualFold(header, "Authorization") {
		token = strings.TrimSpace(r.Header.Get(header))
	}
	if token == "" && options.AllowMock {
		token = bearerToken(r.Header.Get("Authorization"))
	}
	if token == "" {
		return Identity{}, "", fmt.Errorf("missing AgentTeam-verified JWT")
	}
	identity, err := DecodePayload(token)
	if err != nil {
		return Identity{}, "", err
	}
	return identity, token, nil
}

func bearerToken(header string) string {
	fields := strings.Fields(strings.TrimSpace(header))
	if len(fields) == 2 && strings.EqualFold(fields[0], "Bearer") {
		return fields[1]
	}
	return ""
}

func DecodePayload(token string) (Identity, error) {
	parts := strings.Split(token, ".")
	if len(parts) < 2 {
		return Identity{}, fmt.Errorf("invalid JWT: expected header.payload.signature")
	}
	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return Identity{}, fmt.Errorf("invalid JWT payload: %w", err)
	}
	claims := map[string]any{}
	if err := json.Unmarshal(payload, &claims); err != nil {
		return Identity{}, fmt.Errorf("invalid JWT claims: %w", err)
	}
	identity := Identity{RawClaims: claims}
	identity.Subject = stringClaim(claims, "sub")
	identity.Email = stringClaim(claims, "email")
	identity.Name = stringClaim(claims, "name")
	identity.PreferredUsername = stringClaim(claims, "preferred_username")
	identity.Tenant = firstStringClaim(claims, "tenant", "tenant_id", "org", "organization")
	identity.Groups = stringSliceClaim(claims, "groups")
	return identity, nil
}

func stringClaim(claims map[string]any, key string) string {
	value, _ := claims[key].(string)
	return value
}

func firstStringClaim(claims map[string]any, keys ...string) string {
	for _, key := range keys {
		if value := stringClaim(claims, key); value != "" {
			return value
		}
	}
	return ""
}

func stringSliceClaim(claims map[string]any, key string) []string {
	raw, ok := claims[key]
	if !ok {
		return nil
	}
	switch values := raw.(type) {
	case []string:
		return values
	case []any:
		out := make([]string, 0, len(values))
		for _, value := range values {
			if s, ok := value.(string); ok && s != "" {
				out = append(out, s)
			}
		}
		return out
	case string:
		if values == "" {
			return nil
		}
		return strings.Split(values, ",")
	default:
		return nil
	}
}