package oidc

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"net/url"
	"strings"
	"sync"
	"time"
)

// HandoffKind distinguishes registration from MFA-enrollment return envelopes.
type HandoffKind string

const (
	HandoffRegister HandoffKind = "register"
	HandoffEnroll   HandoffKind = "enroll"
)

const defaultHandoffTTL = 10 * time.Minute

// HandoffEnvelope is the signed, expiring state carried to an allow-listed
// registration or enrollment URL. Completing a handoff restarts /authorize
// and never mints a token.
type HandoffEnvelope struct {
	Kind                HandoffKind `json:"kind"`
	ClientID            string      `json:"client_id"`
	RedirectURI         string      `json:"redirect_uri"`
	PKCEChallenge       string      `json:"code_challenge"`
	PKCEChallengeMethod string      `json:"code_challenge_method"`
	State               string      `json:"state"`
	Nonce               string      `json:"nonce,omitempty"`
	Scopes              []string    `json:"scopes,omitempty"`
	TenantHint          string      `json:"tenant_hint,omitempty"`
	JTI                 string      `json:"jti"`
	ExpiresAt           time.Time   `json:"exp"`
}

var (
	errHandoffInvalid = errors.New("invalid handoff")
	errHandoffExpired = errors.New("handoff expired")
	errHandoffReplay  = errors.New("handoff replayed")
)

// HandoffStore signs and atomically consumes registration/enrollment envelopes.
type HandoffStore struct {
	secret []byte
	ttl    time.Duration

	mu       sync.Mutex
	consumed map[string]time.Time
}

// NewHandoffStore returns a store with an ephemeral HMAC key. Envelopes are
// short-lived, so a process restart simply invalidates in-flight handoffs.
func NewHandoffStore() *HandoffStore {
	secret := make([]byte, 32)
	if _, err := rand.Read(secret); err != nil {
		panic("oidc: failed to generate handoff secret: " + err.Error())
	}
	return &HandoffStore{
		secret:   secret,
		ttl:      defaultHandoffTTL,
		consumed: make(map[string]time.Time),
	}
}

// Issue signs a new envelope for the given pending authorization.
func (s *HandoffStore) Issue(kind HandoffKind, ps *PendingState) (string, error) {
	if s == nil {
		return "", errHandoffInvalid
	}
	jti, err := randomID()
	if err != nil {
		return "", err
	}
	env := HandoffEnvelope{
		Kind:                kind,
		ClientID:            ps.ClientID,
		RedirectURI:         ps.RedirectURI,
		PKCEChallenge:       ps.PKCEChallenge,
		PKCEChallengeMethod: ps.PKCEChallengeMethod,
		State:               ps.State,
		Nonce:               ps.Nonce,
		Scopes:              append([]string(nil), ps.Scopes...),
		TenantHint:          ps.TenantHint,
		JTI:                 jti,
		ExpiresAt:           time.Now().Add(s.ttl),
	}
	payload, err := json.Marshal(env)
	if err != nil {
		return "", err
	}
	mac := hmac.New(sha256.New, s.secret)
	mac.Write(payload)
	token := base64.RawURLEncoding.EncodeToString(payload) + "." +
		base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
	return token, nil
}

// Consume verifies the envelope, rejects expiry/tamper/replay, and marks the
// JTI used. The caller must restart /authorize; it must not mint a token.
func (s *HandoffStore) Consume(token string) (*HandoffEnvelope, error) {
	if s == nil {
		return nil, errHandoffInvalid
	}
	payload, err := s.verify(token)
	if err != nil {
		return nil, err
	}
	var env HandoffEnvelope
	if err := json.Unmarshal(payload, &env); err != nil {
		return nil, errHandoffInvalid
	}
	if env.JTI == "" || env.ClientID == "" || env.RedirectURI == "" {
		return nil, errHandoffInvalid
	}
	if time.Now().After(env.ExpiresAt) {
		return nil, errHandoffExpired
	}

	s.mu.Lock()
	defer s.mu.Unlock()
	s.gcLocked()
	if _, used := s.consumed[env.JTI]; used {
		return nil, errHandoffReplay
	}
	s.consumed[env.JTI] = env.ExpiresAt
	return &env, nil
}

func (s *HandoffStore) verify(token string) ([]byte, error) {
	dot := strings.LastIndex(token, ".")
	if dot <= 0 || dot == len(token)-1 {
		return nil, errHandoffInvalid
	}
	payload, err := base64.RawURLEncoding.DecodeString(token[:dot])
	if err != nil {
		return nil, errHandoffInvalid
	}
	sig, err := base64.RawURLEncoding.DecodeString(token[dot+1:])
	if err != nil {
		return nil, errHandoffInvalid
	}
	mac := hmac.New(sha256.New, s.secret)
	mac.Write(payload)
	if !hmac.Equal(mac.Sum(nil), sig) {
		return nil, errHandoffInvalid
	}
	return payload, nil
}

func (s *HandoffStore) gcLocked() {
	now := time.Now()
	for jti, exp := range s.consumed {
		if now.After(exp) {
			delete(s.consumed, jti)
		}
	}
}

func randomID() (string, error) {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(b), nil
}

func appendHandoff(destURL, token string) (string, error) {
	u, err := url.Parse(destURL)
	if err != nil || u.Scheme == "" || u.Host == "" {
		return "", errHandoffInvalid
	}
	q := u.Query()
	q.Set("kc_handoff", token)
	u.RawQuery = q.Encode()
	return u.String(), nil
}