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.
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"
}

// 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.
func EncodeState(secret []byte, provider, next 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,
	}
	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.
func DecodeState(secret []byte, raw string) (provider, next 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, 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")