package openai

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

	"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"
	principalSourceLegacy  = "legacy_bearer_token"
	principalSourceUnknown = "unknown"
)

// 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"
)

// 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{}

// 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
}

// 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 principalFromRequest(r *http.Request, cfg config.EdgeOpenAIConf) (openAIPrincipal, bool) {
	return resolvePrincipal(cfg, r.Header.Get("Authorization"))
}