package ect

import (
	"crypto/ecdsa"
	"crypto/subtle"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"github.com/go-jose/go-jose/v4"
)

// ParsedECT holds a verified or parsed ECT (header + payload).
type ParsedECT struct {
	Header  *jose.Header
	Payload *Payload
	Raw     string // compact JWS
}

// KeyResolver returns the public key for a given kid, or nil if unknown/revoked.
// Returning (nil, nil) means key not found; (nil, err) means lookup error.
type KeyResolver func(kid string) (*ecdsa.PublicKey, error)

// VerifyOptions configures ECT verification per Section 7.
type VerifyOptions struct {
	// VerifierID is the workload identity of the verifier (must be in aud).
	VerifierID string
	// ResolveKey returns the public key for kid. Required.
	ResolveKey KeyResolver
	// ECTStore for DAG validation. If nil, DAG validation is skipped.
	Store ECTStore
	// DAG for DAG validation. Used only if Store != nil.
	DAG DAGConfig
	// Now is the current time; if zero, time.Now() is used.
	Now time.Time
	// IATMaxAge is max age of iat (recommended 15 min). If zero, 15 minutes.
	IATMaxAge time.Duration
	// IATMaxFuture is max clock skew for iat in future (recommended 30 sec). If zero, 30 seconds.
	IATMaxFuture time.Duration
	// JTISeen returns true if jti was already seen (replay). If nil, replay check is skipped.
	JTISeen func(jti string) bool
	// WITSubject if set must equal payload.iss (issuer must match WIT subject for this key).
	WITSubject string
	// ValidateUUIDs when true requires jti and wid (if set) to be UUID format.
	ValidateUUIDs bool
	// MaxParLength caps par length (0 = no limit). Applied before DAG; DAG may also enforce via DAG.MaxParLength.
	MaxParLength int
	// LogVerify if set is called after verification with jti and any error (for observability).
	LogVerify func(jti string, err error)
}

// DefaultVerifyOptions returns recommended defaults.
func DefaultVerifyOptions() VerifyOptions {
	return VerifyOptions{
		IATMaxAge:     15 * time.Minute,
		IATMaxFuture:  30 * time.Second,
		DAG:           DefaultDAGConfig(),
	}
}

// Parse parses compact JWS and returns header + payload without cryptographic verification.
// Use Verify for full signature and claim validation.
func Parse(compact string) (*ParsedECT, error) {
	jws, err := jose.ParseSigned(compact, []jose.SignatureAlgorithm{jose.ES256})
	if err != nil {
		return nil, err
	}
	if len(jws.Signatures) != 1 {
		return nil, errors.New("ect: expected single signature")
	}
	sig := jws.Signatures[0]
	payloadBytes := jws.UnsafePayloadWithoutVerification()
	if len(payloadBytes) == 0 {
		return nil, errors.New("ect: empty payload")
	}
	var p Payload
	if err := json.Unmarshal(payloadBytes, &p); err != nil {
		return nil, err
	}
	return &ParsedECT{
		Header:  &sig.Header,
		Payload: &p,
		Raw:     compact,
	}, nil
}

// Verify performs full Section 7 verification and optional DAG validation.
func Verify(compact string, opts VerifyOptions) (parsed *ParsedECT, err error) {
	var logJti string
	defer func() {
		if opts.LogVerify != nil {
			opts.LogVerify(logJti, err)
		}
	}()

	jws, err := jose.ParseSigned(compact, []jose.SignatureAlgorithm{jose.ES256})
	if err != nil {
		return nil, err
	}
	if len(jws.Signatures) != 1 {
		return nil, errors.New("ect: expected single signature")
	}
	sig := jws.Signatures[0]
	header := &sig.Header

	// 2. typ must be wimse-exec+jwt (constant-time compare)
	typ, _ := header.ExtraHeaders["typ"].(string)
	if typ == "" {
		typ, _ = header.ExtraHeaders[jose.HeaderType].(string)
	}
	if subtle.ConstantTimeCompare([]byte(typ), []byte(ECTType)) != 1 {
		return nil, ErrInvalidTyp
	}

	// 3. alg not none and not symmetric
	if header.Algorithm == "none" || header.Algorithm == "HS256" || header.Algorithm == "HS384" || header.Algorithm == "HS512" {
		return nil, ErrProhibitedAlg
	}

	// 4. kid references known key
	kid := header.KeyID
	if kid == "" {
		return nil, ErrMissingKid
	}
	pub, err := opts.ResolveKey(kid)
	if err != nil {
		return nil, fmt.Errorf("ect: key resolution: %w", err)
	}
	if pub == nil {
		return nil, ErrUnknownKey
	}

	// 5. Verify JWS signature
	payloadBytes, err := jws.Verify(&jose.JSONWebKey{Key: pub, KeyID: kid})
	if err != nil {
		return nil, fmt.Errorf("ect: invalid signature: %w", err)
	}

	var p Payload
	if err := json.Unmarshal(payloadBytes, &p); err != nil {
		return nil, err
	}
	logJti = p.Jti

	if err := ValidateExt(p.Ext); err != nil {
		return nil, err
	}

	// 6. Key revocation — caller's ResolveKey can return nil for revoked
	// 7. alg match WIT — we don't have WIT; optional WITSubject check below
	// 8. iss matches WIT subject
	if opts.WITSubject != "" && p.Iss != opts.WITSubject {
		return nil, ErrWITSubjectMismatch
	}

	// 9. aud contains verifier
	if opts.VerifierID != "" && !p.ContainsAudience(opts.VerifierID) {
		return nil, ErrAudienceMismatch
	}

	// 10. exp not expired
	now := opts.Now
	if now.IsZero() {
		now = time.Now()
	}
	if now.Unix() > p.Exp {
		return nil, ErrExpired
	}

	// 11. iat not too far past/future
	if opts.IATMaxAge == 0 {
		opts.IATMaxAge = 15 * time.Minute
	}
	if opts.IATMaxFuture == 0 {
		opts.IATMaxFuture = 30 * time.Second
	}
	if now.Unix()-p.Iat > int64(opts.IATMaxAge.Seconds()) {
		return nil, ErrIATTooOld
	}
	if p.Iat > now.Unix()+int64(opts.IATMaxFuture.Seconds()) {
		return nil, ErrIATInFuture
	}

	// 12. Required claims present (jti, exec_act, par)
	if p.Jti == "" || p.ExecAct == "" {
		return nil, ErrMissingClaims
	}
	if p.Par == nil {
		p.Par = []string{}
	}
	if opts.MaxParLength > 0 && len(p.Par) > opts.MaxParLength {
		return nil, ErrParLength
	}
	if opts.ValidateUUIDs {
		if !ValidUUID(p.Jti) {
			return nil, ErrInvalidJTI
		}
		if p.Wid != "" && !ValidUUID(p.Wid) {
			return nil, ErrInvalidWID
		}
	}
	if p.InpHash != "" {
		if err := ValidateHashFormat(p.InpHash); err != nil {
			return nil, err
		}
	}
	if p.OutHash != "" {
		if err := ValidateHashFormat(p.OutHash); err != nil {
			return nil, err
		}
	}

	// 13. If pol or pol_decision present, both must be present and pol_decision in registry
	if p.Pol != "" || p.PolDecision != "" {
		if p.Pol == "" || p.PolDecision == "" {
			return nil, ErrPolPolDecisionPair
		}
		if !ValidPolDecision(p.PolDecision) {
			return nil, ErrInvalidPolDecision
		}
	}

	// 14. DAG validation
	if opts.Store != nil {
		if err := ValidateDAG(&p, opts.Store, opts.DAG); err != nil {
			return nil, err
		}
	}

	// 15. Replay (jti seen)
	if opts.JTISeen != nil && opts.JTISeen(p.Jti) {
		return nil, ErrReplay
	}

	return &ParsedECT{Header: header, Payload: &p, Raw: compact}, nil
}