package oidcrp

import (
	"bytes"
	"cmp"
	"context"
	"crypto/hkdf"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"slices"
	"strings"
	"time"

	"github.com/coreos/go-oidc/v3/oidc"
	"github.com/dexidp/dex/server/tokens"
	"golang.org/x/oauth2"

	"tug/internal/platform/config"
	"tug/internal/platform/spent"
	transporthttp "tug/internal/transport/http"
)

const (
	TxnCookie      = "__Host-authsvc_txn"
	SessionCookie  = "__Host-authsvc_sess"
	SecretLength   = 32
	MaxTxnTTL      = 10 * time.Minute
	maxTxnLimit    = 16
	MaxReturnBytes = 2048
)

var Scopes = []string{oidc.ScopeOpenID, "profile", "groups", tokens.ScopeFederatedID}

var (
	ErrUnknownHost   = errors.New("oidcrp: the host has no client")
	ErrReturn        = errors.New("oidcrp: the return path is not an absolute path")
	ErrNoTxn         = errors.New("oidcrp: no transaction cookie")
	ErrTxnMalformed  = errors.New("oidcrp: the transaction does not decode")
	ErrTxnExpired    = errors.New("oidcrp: the transaction expired")
	ErrTxnHost       = errors.New("oidcrp: the transaction belongs to another host")
	ErrState         = errors.New("oidcrp: state does not match the transaction")
	ErrTxnSpent      = errors.New("oidcrp: the transaction was already spent")
	ErrAuthorization = errors.New("oidcrp: the provider returned an error")
	ErrExchange      = errors.New("oidcrp: the code exchange failed")
	ErrIDToken       = errors.New("oidcrp: the ID token did not verify")
	ErrNonce         = errors.New("oidcrp: nonce does not match the transaction")
	ErrSubject       = errors.New("oidcrp: the subject does not match the federated user id")
	ErrUser          = errors.New("oidcrp: the user is missing or disabled")
	ErrNoSession     = errors.New("oidcrp: no session")
	ErrSessions      = errors.New("oidcrp: the session store failed")
)

type Session struct {
	UserID  string
	Expires time.Time
}

type Sessions interface {
	Current(ctx context.Context, secret []byte, host string) (Session, error)
	OpenAt(ctx context.Context, userID, host, clientID string, epoch int64, browser string) ([]byte, Session, error)
	Revoke(ctx context.Context, secret []byte, host string) error
	Authenticated(ctx context.Context, secret []byte, host, clientID string) error
}

type Client struct {
	ID           string
	Secret       config.Secret
	Origins      []string
	CallbackPath string
}

type Config struct {
	Issuer      string
	IDP         http.Handler
	Hosts       []string
	MasterKey   config.Secret
	Clients     []Client
	ConnectorID string
	Sessions    Sessions
	TxnTTL      time.Duration
	TxnLimit    int
	Now         func() time.Time
	Refuse      func(w http.ResponseWriter, r *http.Request, status int, err error)
	Rotate      func(w http.ResponseWriter)
	CodeEpoch   func(ctx context.Context, code string) (string, int64, string, error)
}

type txn struct {
	State    string `json:"s"`
	Nonce    string `json:"n"`
	Verifier string `json:"v"`
	Host     string `json:"h"`
	Return   string `json:"r"`
	Expiry   int64  `json:"x"`
}

type target struct {
	client Client
	origin string
}

type RP struct {
	cfg       Config
	now       func() time.Time
	ttl       time.Duration
	client    *http.Client
	endpoint  oauth2.Endpoint
	byHost    map[string]target
	verifiers map[string]*oidc.IDTokenVerifier
	txn       *Sealer
	sess      *Sealer
	login     *Sealer
	proofKey  []byte
	spent     *spent.Set
}

func New(ctx context.Context, cfg Config) (*RP, error) {
	iss, err := url.Parse(cfg.Issuer)
	if err != nil || iss.Host == "" || (iss.Scheme != "https" && iss.Scheme != "http") {
		return nil, fmt.Errorf("oidcrp: issuer %q is not scheme://host", cfg.Issuer)
	}
	if cfg.IDP == nil || cfg.Sessions == nil || cfg.Refuse == nil || cfg.Rotate == nil || cfg.ConnectorID == "" {
		return nil, errors.New("oidcrp: the provider handler, the sessions, the refusal page, the CSRF rotation and the connector id are required")
	}
	if cfg.CodeEpoch == nil {
		return nil, errors.New("oidcrp: the code epoch reader is required")
	}
	ttl := cfg.TxnTTL
	if ttl == 0 {
		ttl = MaxTxnTTL
	}
	if ttl < 0 || ttl > MaxTxnTTL {
		return nil, errors.New("oidcrp: the transaction lifetime is at most ten minutes")
	}
	if cfg.TxnLimit < 0 || cfg.TxnLimit > maxTxnLimit {
		return nil, fmt.Errorf("oidcrp: the transaction limit is between 0 and %d", maxTxnLimit)
	}
	now := cfg.Now
	if now == nil {
		now = time.Now
	}
	key := cfg.MasterKey.Bytes()
	txnSealer, err := NewSealer(key, TxnInfo)
	if err != nil {
		return nil, err
	}
	sessSealer, err := NewSealer(key, SessionInfo)
	if err != nil {
		return nil, err
	}
	loginSealer, err := NewSealer(key, LoginInfo)
	if err != nil {
		return nil, err
	}
	proofKey, err := hkdf.Key(sha256.New, key, nil, LoginProofInfo, LoginProofLength)
	if err != nil {
		return nil, err
	}
	byHost := map[string]target{}
	for _, c := range cfg.Clients {
		if c.ID == "" || len(c.Secret.Bytes()) == 0 || !strings.HasPrefix(c.CallbackPath, "/") || len(c.Origins) == 0 {
			return nil, fmt.Errorf("oidcrp: client %q needs a secret, a callback path and an origin", c.ID)
		}
		for _, raw := range c.Origins {
			host, err := transporthttp.ParseOrigin(raw, iss.Scheme, cfg.Hosts...)
			if err != nil {
				return nil, fmt.Errorf("oidcrp: client %q: %w", c.ID, err)
			}
			if _, dup := byHost[host]; dup {
				return nil, fmt.Errorf("oidcrp: host %q has two origins: %w", host, transporthttp.ErrOrigin)
			}
			byHost[host] = target{client: c, origin: raw}
		}
	}
	client := &http.Client{
		Transport:     InProcess(iss.Host, cfg.IDP),
		Timeout:       15 * time.Second,
		CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse },
	}
	provider, err := oidc.NewProvider(oidc.ClientContext(context.WithoutCancel(ctx), client), cfg.Issuer)
	if err != nil {
		return nil, fmt.Errorf("oidcrp: discovery: %w", err)
	}
	endpoint := provider.Endpoint()
	endpoint.AuthStyle = oauth2.AuthStyleInHeader
	rp := &RP{
		cfg: cfg, now: now, ttl: ttl, client: client, endpoint: endpoint,
		byHost: byHost, verifiers: map[string]*oidc.IDTokenVerifier{},
		txn: txnSealer, sess: sessSealer, login: loginSealer, proofKey: proofKey, spent: spent.New(),
	}
	for _, c := range cfg.Clients {
		rp.verifiers[c.ID] = provider.Verifier(&oidc.Config{ClientID: c.ID, Now: now})
	}
	return rp, nil
}

func (rp *RP) oauth(tg target) *oauth2.Config {
	return &oauth2.Config{
		ClientID: tg.client.ID, ClientSecret: string(tg.client.Secret.Bytes()), Endpoint: rp.endpoint,
		RedirectURL: tg.origin + tg.client.CallbackPath, Scopes: Scopes,
	}
}

func ValidReturn(p string) error {
	if !strings.HasPrefix(p, "/") || strings.HasPrefix(p, "//") || strings.HasPrefix(p, "/\\") {
		return ErrReturn
	}
	u, err := url.Parse(p)
	if err != nil || u.Scheme != "" || u.Host != "" || u.User != nil || u.Opaque != "" {
		return ErrReturn
	}
	if b, err := json.Marshal(p); err != nil || len(b)-2 > MaxReturnBytes {
		return ErrReturn
	}
	return nil
}

func ReturnOrRoot(p string) string {
	if ValidReturn(p) != nil {
		return "/"
	}
	return p
}

func random() string {
	b := make([]byte, SecretLength)
	rand.Read(b)
	return base64.RawURLEncoding.EncodeToString(b)
}

func cookie(name, value string, maxAge int) *http.Cookie {
	return &http.Cookie{Name: name, Value: value, Path: "/", MaxAge: maxAge, Secure: true, HttpOnly: true, SameSite: http.SameSiteLaxMode}
}

func (rp *RP) Begin(w http.ResponseWriter, r *http.Request, returnPath string) {
	w.Header().Set("Cache-Control", "no-store")
	host := transporthttp.CanonicalHost(r.Host)
	tg, ok := rp.byHost[host]
	if !ok {
		rp.cfg.Refuse(w, r, http.StatusNotFound, ErrUnknownHost)
		return
	}
	if err := ValidReturn(returnPath); err != nil {
		rp.cfg.Refuse(w, r, http.StatusBadRequest, err)
		return
	}
	t := txn{
		State: random(), Nonce: random(), Verifier: oauth2.GenerateVerifier(), Host: host,
		Return: returnPath, Expiry: rp.now().Add(rp.ttl).Unix(),
	}
	b, err := json.Marshal(t)
	if err != nil {
		rp.cfg.Refuse(w, r, http.StatusInternalServerError, err)
		return
	}
	name := rp.txnName(t.State)
	if rp.cfg.TxnLimit > 0 {
		rp.evictTxns(w, r, host)
	}
	http.SetCookie(w, cookie(name, rp.txn.Seal(name, b), int(rp.ttl/time.Second)))
	http.Redirect(w, r, rp.oauth(tg).AuthCodeURL(t.State, oauth2.S256ChallengeOption(t.Verifier), oidc.Nonce(t.Nonce)), http.StatusFound)
}

func txnCookieName(state string) string {
	return cookieName(TxnCookie, state)
}

func cookieName(prefix, state string) string {
	sum := sha256.Sum256([]byte(state))
	return prefix + "_" + base64.RawURLEncoding.EncodeToString(sum[:])[:11]
}

func (rp *RP) txnName(state string) string {
	if rp.cfg.TxnLimit == 0 {
		return TxnCookie
	}
	return txnCookieName(state)
}

func (rp *RP) evictTxns(w http.ResponseWriter, r *http.Request, host string) {
	evictCookies(w, r, TxnCookie+"_", rp.cfg.TxnLimit, rp.now(), func(c *http.Cookie) (time.Time, bool) {
		t, err := rp.decodeTxn(c.Name, c.Value)
		return time.Unix(t.Expiry, 0), err == nil && t.Host == host
	})
}

func evictCookies(w http.ResponseWriter, r *http.Request, prefix string, limit int, now time.Time, decode func(*http.Cookie) (time.Time, bool)) {
	if limit <= 0 {
		return
	}
	live := liveCookies(w, r, prefix, now, decode)
	for len(live) >= limit {
		http.SetCookie(w, cookie(live[0].name, "", -1))
		live = live[1:]
	}
}

type heldCookie struct {
	name   string
	expiry int64
}

func liveCookies(w http.ResponseWriter, r *http.Request, prefix string, now time.Time, decode func(*http.Cookie) (time.Time, bool)) []heldCookie {
	var live []heldCookie
	for _, ck := range r.Cookies() {
		if !strings.HasPrefix(ck.Name, prefix) {
			continue
		}
		expiry, ok := decode(ck)
		if !ok || !now.Before(expiry) {
			http.SetCookie(w, cookie(ck.Name, "", -1))
			continue
		}
		live = append(live, heldCookie{ck.Name, expiry.UnixMicro()})
	}
	slices.SortStableFunc(live, func(a, b heldCookie) int { return cmp.Compare(a.expiry, b.expiry) })
	return live
}

func (rp *RP) openTxn(r *http.Request, name string) (txn, error) {
	ck, err := r.Cookie(name)
	if err != nil {
		return txn{}, ErrNoTxn
	}
	return rp.decodeTxn(name, ck.Value)
}

func (rp *RP) decodeTxn(name, value string) (txn, error) {
	raw, err := rp.txn.Open(name, value)
	if err != nil {
		return txn{}, err
	}
	dec := json.NewDecoder(bytes.NewReader(raw))
	dec.DisallowUnknownFields()
	var t txn
	if err := dec.Decode(&t); err != nil || dec.More() || t.State == "" || t.Nonce == "" || t.Verifier == "" || ValidReturn(t.Return) != nil {
		return txn{}, ErrTxnMalformed
	}
	return t, nil
}

func CheckSubject(sub, userID, connectorID string) error {
	if userID == "" {
		return ErrSubject
	}
	want, err := tokens.GenSubject(userID, connectorID)
	if err != nil {
		return fmt.Errorf("%w: %w", ErrSubject, err)
	}
	if subtle.ConstantTimeCompare([]byte(sub), []byte(want)) != 1 {
		return ErrSubject
	}
	return nil
}

type federated struct {
	Claims struct {
		UserID string `json:"user_id"`
	} `json:"federated_claims"`
}

func (rp *RP) Callback(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Cache-Control", "no-store")
	host := transporthttp.CanonicalHost(r.Host)
	tg, ok := rp.byHost[host]
	if !ok {
		rp.cfg.Refuse(w, r, http.StatusNotFound, ErrUnknownHost)
		return
	}
	q, err := transporthttp.Query(r)
	if err != nil {
		rp.cfg.Refuse(w, r, http.StatusBadRequest, err)
		return
	}
	state := q.Get("state")
	name := rp.txnName(state)
	t, err := rp.openTxn(r, name)
	if err != nil {
		rp.cfg.Refuse(w, r, http.StatusBadRequest, err)
		return
	}
	expiry := time.Unix(t.Expiry, 0)
	if !rp.now().Before(expiry) {
		http.SetCookie(w, cookie(name, "", -1))
		rp.cfg.Refuse(w, r, http.StatusBadRequest, ErrTxnExpired)
		return
	}
	if t.Host != host {
		rp.cfg.Refuse(w, r, http.StatusBadRequest, ErrTxnHost)
		return
	}
	if subtle.ConstantTimeCompare([]byte(state), []byte(t.State)) != 1 {
		rp.cfg.Refuse(w, r, http.StatusBadRequest, ErrState)
		return
	}
	http.SetCookie(w, cookie(name, "", -1))
	if !rp.spent.Spend(t.State, expiry, rp.now()) {
		rp.cfg.Refuse(w, r, http.StatusBadRequest, ErrTxnSpent)
		return
	}
	if q.Get("error") != "" || q.Get("code") == "" {
		rp.cfg.Refuse(w, r, http.StatusBadRequest, ErrAuthorization)
		return
	}
	codeUser, epoch, browser, err := rp.cfg.CodeEpoch(r.Context(), q.Get("code"))
	if err != nil {
		rp.cfg.Refuse(w, r, http.StatusBadGateway, fmt.Errorf("%w: %w", ErrExchange, err))
		return
	}
	ctx := context.WithValue(r.Context(), oauth2.HTTPClient, rp.client)
	tok, err := rp.oauth(tg).Exchange(ctx, q.Get("code"), oauth2.VerifierOption(t.Verifier))
	if err != nil {
		rp.cfg.Refuse(w, r, http.StatusBadGateway, fmt.Errorf("%w: %w", ErrExchange, err))
		return
	}
	rawID, _ := tok.Extra("id_token").(string)
	if rawID == "" {
		rp.cfg.Refuse(w, r, http.StatusBadGateway, ErrExchange)
		return
	}
	idt, err := rp.verifiers[tg.client.ID].Verify(oidc.ClientContext(r.Context(), rp.client), rawID)
	if err != nil {
		rp.cfg.Refuse(w, r, http.StatusUnauthorized, fmt.Errorf("%w: %w", ErrIDToken, err))
		return
	}
	if subtle.ConstantTimeCompare([]byte(idt.Nonce), []byte(t.Nonce)) != 1 {
		rp.cfg.Refuse(w, r, http.StatusUnauthorized, ErrNonce)
		return
	}
	var fc federated
	if err := idt.Claims(&fc); err != nil {
		rp.cfg.Refuse(w, r, http.StatusUnauthorized, ErrSubject)
		return
	}
	if err := CheckSubject(idt.Subject, fc.Claims.UserID, rp.cfg.ConnectorID); err != nil {
		rp.cfg.Refuse(w, r, http.StatusUnauthorized, err)
		return
	}
	if codeUser != fc.Claims.UserID {
		rp.cfg.Refuse(w, r, http.StatusUnauthorized, fmt.Errorf("%w: the code was issued to another user", ErrSubject))
		return
	}
	secret, opened, err := rp.establish(r, host, tg.client.ID, fc.Claims.UserID, epoch, browser)
	switch {
	case errors.Is(err, ErrUser):
		rp.cfg.Refuse(w, r, http.StatusForbidden, err)
		return
	case err != nil:
		rp.cfg.Refuse(w, r, http.StatusInternalServerError, err)
		return
	}
	if opened == nil {
		if err := rp.cfg.Sessions.Authenticated(r.Context(), secret, host, tg.client.ID); err != nil {
			rp.cfg.Refuse(w, r, http.StatusInternalServerError, fmt.Errorf("%w: %w", ErrSessions, err))
			return
		}
	} else {
		rp.SetSession(w, secret, *opened)
		rp.cfg.Rotate(w)
	}
	http.Redirect(w, r, tg.origin+t.Return, http.StatusSeeOther)
}

func (rp *RP) establish(r *http.Request, host, clientID, userID string, epoch int64, browser string) ([]byte, *Session, error) {
	ctx := r.Context()
	if secret, ok := rp.SessionSecret(r); ok {
		s, err := rp.cfg.Sessions.Current(ctx, secret, host)
		switch {
		case err == nil && s.UserID == userID:
			return secret, nil, nil
		case err == nil:
			if err := rp.cfg.Sessions.Revoke(ctx, secret, host); err != nil {
				return nil, nil, fmt.Errorf("%w: %w", ErrSessions, err)
			}
		case !errors.Is(err, ErrNoSession):
			return nil, nil, fmt.Errorf("%w: %w", ErrSessions, err)
		}
	}
	secret, s, err := rp.cfg.Sessions.OpenAt(ctx, userID, host, clientID, epoch, browser)
	if err != nil {
		return nil, nil, fmt.Errorf("%w: %w", ErrSessions, err)
	}
	return secret, &s, nil
}

func (rp *RP) SessionSecret(r *http.Request) ([]byte, bool) {
	ck, err := r.Cookie(SessionCookie)
	if err != nil {
		return nil, false
	}
	secret, err := rp.sess.Open(SessionCookie, ck.Value)
	if err != nil || len(secret) != SecretLength {
		return nil, false
	}
	return secret, true
}

func (rp *RP) SetSession(w http.ResponseWriter, secret []byte, s Session) {
	maxAge := int(s.Expires.Sub(rp.now()) / time.Second)
	if maxAge < 1 {
		maxAge = 1
	}
	http.SetCookie(w, cookie(SessionCookie, rp.sess.Seal(SessionCookie, secret), maxAge))
}

func (rp *RP) ClearSession(w http.ResponseWriter) {
	http.SetCookie(w, cookie(SessionCookie, "", -1))
}