package oauth

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"
)

// statePayload is the JSON shape we round-trip through the OAuth state
// parameter. The cookie holds the same encoded blob; on callback we compare
// the two byte-for-byte before trusting either.
//
// `Next` is the post-login redirect path on the frontend. We sign it (along
// with Nonce + IssuedAt) under HMAC so a tampered ?next= can't redirect the
// user to an arbitrary URL after sign-in.
//
// `LinkUserID` is set only by the /link/start endpoint — it pins the OAuth
// callback to the user who initiated the link. An empty value means the
// callback should run the normal sign-in / sign-up flow instead.
type statePayload struct {
	Nonce      string `json:"n"`            // 16 random bytes, base64-urlencoded
	Next       string `json:"u"`            // safe path, validated by handler
	IssuedAt   int64  `json:"t"`            // unix seconds
	Provider   string `json:"p"`            // "google" | "facebook"
	LinkUserID string `json:"l,omitempty"`  // user id to link the provider to, empty for sign-in flow
}

// stateMaxAge bounds how long a /start may be in flight. Ten minutes is
// generous — typical OAuth round-trips finish in seconds. Longer than this
// and we discard the state to keep the replay window small.
const stateMaxAge = 10 * time.Minute

// EncodeState signs and packs a fresh state value bound to the given provider
// and post-login redirect path. The returned string is what goes into both
// the `state` query param and the cookie.
//
// `linkUserID` pins the callback to a specific user (link-to-current-user
// flow). Pass "" for the normal sign-in flow.
func EncodeState(secret []byte, provider, next, linkUserID string) (string, error) {
	nonceB := make([]byte, 16)
	if _, err := rand.Read(nonceB); err != nil {
		return "", fmt.Errorf("state nonce: %w", err)
	}
	p := statePayload{
		Nonce:      base64.RawURLEncoding.EncodeToString(nonceB),
		Next:       next,
		IssuedAt:   time.Now().Unix(),
		Provider:   provider,
		LinkUserID: linkUserID,
	}
	body, err := json.Marshal(p)
	if err != nil {
		return "", err
	}
	bodyEnc := base64.RawURLEncoding.EncodeToString(body)
	mac := hmac.New(sha256.New, secret)
	mac.Write([]byte(bodyEnc))
	sig := base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
	return bodyEnc + "." + sig, nil
}

// DecodeState verifies the signature, checks the freshness window, and
// returns the embedded payload. Any failure surfaces as ErrInvalidState so
// the handler maps it to a single error path.
//
// `linkUserID` is empty for sign-in flows and set for link-to-current-user
// flows — callers should branch on it.
func DecodeState(secret []byte, raw string) (provider, next, linkUserID string, err error) {
	parts := strings.SplitN(raw, ".", 2)
	if len(parts) != 2 {
		return "", "", "", ErrInvalidState
	}
	bodyEnc, sigEnc := parts[0], parts[1]
	mac := hmac.New(sha256.New, secret)
	mac.Write([]byte(bodyEnc))
	expected := base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
	if !hmac.Equal([]byte(expected), []byte(sigEnc)) {
		return "", "", "", ErrInvalidState
	}
	body, err := base64.RawURLEncoding.DecodeString(bodyEnc)
	if err != nil {
		return "", "", "", ErrInvalidState
	}
	var p statePayload
	if err := json.Unmarshal(body, &p); err != nil {
		return "", "", "", ErrInvalidState
	}
	if time.Since(time.Unix(p.IssuedAt, 0)) > stateMaxAge {
		return "", "", "", ErrInvalidState
	}
	return p.Provider, p.Next, p.LinkUserID, nil
}

// ErrInvalidState is returned whenever signature or freshness checks fail.
// We deliberately return one error so handlers don't leak whether a request
// was tampered or just stale.
var ErrInvalidState = errors.New("invalid or expired oauth state")