package openai

import (
	"context"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"net/http"
	"strings"

	"iop/apps/edge/internal/authprojection"
	"iop/packages/go/config"
)

// Principal source tags. These label how an openAIPrincipal was resolved and
// are surfaced in dispatch metadata for usage/audit purposes.
const (
	principalSourceToken      = "principal_token"
	principalSourceProjection = "control_plane_projection"
	principalSourceLegacy     = "legacy_bearer_token"
	principalSourceUnknown    = "unknown"
)

type credentialMode uint8

const (
	credentialModeLegacy credentialMode = iota
	credentialModeManaged
)

// Dispatch metadata keys carrying the resolved principal. These are set from
// the authenticated request context only; metadata.user from the caller's
// request body is never used as an identity source.
const (
	principalMetaRef       = "iop_principal_ref"
	principalMetaAlias     = "iop_principal_alias"
	principalMetaToken     = "iop_token_ref"
	principalMetaSource    = "iop_principal_source"
	routeMetaID            = "iop_route_id"
	routeMetaRevision      = "iop_route_revision"
	credentialSlotMeta     = "iop_credential_slot_ref"
	credentialRevisionMeta = "iop_credential_revision"
)

// openAIPrincipal is the caller identity resolved from an OpenAI-compatible
// bearer token via IOP bearer-token-operation mapping. It never carries the
// raw token value.
type openAIPrincipal struct {
	PrincipalRef   string
	PrincipalAlias string
	TokenRef       string
	Source         string
}

type principalContextKey struct{}
type projectionViewContextKey struct{}

// withPrincipal returns a context carrying the resolved caller principal.
func withPrincipal(ctx context.Context, p openAIPrincipal) context.Context {
	return context.WithValue(ctx, principalContextKey{}, p)
}

// principalFromContext returns the caller principal stored by withAuth, if
// any.
func principalFromContext(ctx context.Context) (openAIPrincipal, bool) {
	p, ok := ctx.Value(principalContextKey{}).(openAIPrincipal)
	return p, ok
}

func withAuthenticatedProjectionView(ctx context.Context, view authprojection.AuthenticatedView) context.Context {
	return context.WithValue(ctx, projectionViewContextKey{}, view)
}

func authenticatedProjectionViewFromContext(ctx context.Context) (authprojection.AuthenticatedView, bool) {
	view, ok := ctx.Value(projectionViewContextKey{}).(authprojection.AuthenticatedView)
	return view, ok && view.State == authprojection.StateFresh
}

// principalMetadata renders the request context's principal as dispatch
// metadata. It returns nil when no principal was resolved for the context.
func principalMetadata(ctx context.Context) map[string]string {
	p, ok := principalFromContext(ctx)
	if !ok {
		return nil
	}
	meta := map[string]string{principalMetaSource: p.Source}
	if p.PrincipalRef != "" {
		meta[principalMetaRef] = p.PrincipalRef
	}
	if p.PrincipalAlias != "" {
		meta[principalMetaAlias] = p.PrincipalAlias
	}
	if p.TokenRef != "" {
		meta[principalMetaToken] = p.TokenRef
	}
	return meta
}

// bearerTokenFromHeader extracts the raw bearer token value from an
// Authorization header, or "" if the header is missing or malformed.
func bearerTokenFromHeader(authHeader string) string {
	const prefix = "Bearer "
	if !strings.HasPrefix(authHeader, prefix) {
		return ""
	}
	return strings.TrimPrefix(authHeader, prefix)
}

// resolvePrincipal authenticates an OpenAI-compatible bearer token against
// cfg.PrincipalTokens (hash mapping, checked first) and cfg.BearerToken
// (legacy fallback). It returns the resolved principal and whether the
// caller is authenticated. When neither principal_tokens nor bearer_token is
// configured, the endpoint is open: authenticated=true with an
// unknown-source principal, matching the pre-existing no-auth default.
func resolvePrincipal(cfg config.EdgeOpenAIConf, authHeader string) (openAIPrincipal, bool) {
	hasPrincipalTokens := len(cfg.PrincipalTokens) > 0
	hasLegacyToken := strings.TrimSpace(cfg.BearerToken) != ""
	if !hasPrincipalTokens && !hasLegacyToken {
		return openAIPrincipal{Source: principalSourceUnknown}, true
	}

	token := bearerTokenFromHeader(authHeader)
	if token == "" {
		return openAIPrincipal{}, false
	}

	if hasPrincipalTokens {
		if p, ok := matchPrincipalToken(cfg.PrincipalTokens, token); ok {
			return p, true
		}
	}

	if hasLegacyToken {
		expected := []byte(cfg.BearerToken)
		if subtle.ConstantTimeCompare([]byte(token), expected) == 1 {
			return openAIPrincipal{Source: principalSourceLegacy}, true
		}
	}

	return openAIPrincipal{}, false
}

// matchPrincipalToken hashes token and compares it against each configured
// entry's stored hash using a constant-time comparison per candidate.
func matchPrincipalToken(tokens []config.OpenAIPrincipalTokenConf, token string) (openAIPrincipal, bool) {
	sum := sha256.Sum256([]byte(token))
	hash := hex.EncodeToString(sum[:])
	for _, t := range tokens {
		configured := strings.ToLower(strings.TrimSpace(t.TokenHashSHA256))
		if subtle.ConstantTimeCompare([]byte(hash), []byte(configured)) == 1 {
			return openAIPrincipal{
				PrincipalRef:   t.PrincipalRef,
				PrincipalAlias: t.PrincipalAlias,
				TokenRef:       t.TokenRef,
				Source:         principalSourceToken,
			}, true
		}
	}
	return openAIPrincipal{}, false
}

// principalFromRequest resolves r's caller principal from its Authorization
// header against cfg.
func principalTokenFromRequest(r *http.Request) (string, bool) {
	authorization := strings.TrimSpace(r.Header.Get("Authorization"))
	bearer := bearerTokenFromHeader(authorization)
	apiKey := ""
	if isAnthropicRequest(r) {
		apiKey = strings.TrimSpace(r.Header.Get("X-Api-Key"))
	} else if isGeminiRequest(r) {
		apiKey = strings.TrimSpace(r.Header.Get("X-Goog-Api-Key"))
	}
	if authorization != "" && bearer == "" {
		return "", false
	}
	if bearer != "" && apiKey != "" {
		if subtle.ConstantTimeCompare([]byte(bearer), []byte(apiKey)) != 1 {
			return "", false
		}
	}
	token := bearer
	if token == "" {
		token = apiKey
	}
	return token, true
}

// principalFromRequest resolves r's caller principal from its supported
// surface-specific headers against the legacy static configuration.
func principalFromRequest(r *http.Request, cfg config.EdgeOpenAIConf) (openAIPrincipal, bool) {
	token, ok := principalTokenFromRequest(r)
	if !ok {
		return openAIPrincipal{}, false
	}
	if token == "" {
		return resolvePrincipal(cfg, "")
	}
	return resolvePrincipal(cfg, "Bearer "+token)
}

// authenticatePrincipal selects exactly one authentication source from the
// explicit credential mode. Managed mode requires a fresh shared projection;
// unavailable, unmanaged, expired, and revoked projection states fail closed.
func (s *Server) authenticatePrincipal(r *http.Request) (openAIPrincipal, authprojection.AuthenticatedView, bool) {
	token, headersOK := principalTokenFromRequest(r)
	if s.managedCredentialPlane() {
		projection := s.PrincipalProjection()
		if !headersOK || token == "" || projection == nil {
			return openAIPrincipal{}, authprojection.AuthenticatedView{}, false
		}
		digest := sha256.Sum256([]byte(token))
		view, matched := projection.AuthenticatedView(digest)
		if view.State != authprojection.StateFresh || !matched {
			return openAIPrincipal{}, view, false
		}
		return openAIPrincipal{
			PrincipalRef: view.Principal.PrincipalRef, PrincipalAlias: view.Principal.PrincipalAlias,
			TokenRef: view.Principal.TokenRef, Source: principalSourceProjection,
		}, view, true
	}
	if !headersOK {
		return openAIPrincipal{}, authprojection.AuthenticatedView{}, false
	}
	if token == "" {
		p, ok := resolvePrincipal(s.cfg, "")
		return p, authprojection.AuthenticatedView{}, ok
	}
	p, ok := resolvePrincipal(s.cfg, "Bearer "+token)
	return p, authprojection.AuthenticatedView{}, ok
}
