package turnstile

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/caddyserver/caddy/v2/modules/caddyhttp"
)

const cookieVersion = "2"

type cookieSigner struct {
	secret []byte
	name   string
	ttl    time.Duration
	secure bool
}

func newCookieSigner(secret, name string, ttl time.Duration) (*cookieSigner, error) {
	if secret == "" {
		return nil, fmt.Errorf("cookie_secret is required")
	}
	if name == "" {
		name = "__turnstile_pass"
	}
	if ttl < time.Minute {
		ttl = time.Hour
	}
	return &cookieSigner{
		secret: []byte(secret),
		name:   name,
		ttl:    ttl,
		secure: true,
	}, nil
}

func (c *cookieSigner) set(w http.ResponseWriter, r *http.Request) error {
	nonce := make([]byte, 16)
	if _, err := rand.Read(nonce); err != nil {
		return err
	}
	exp := time.Now().Add(c.ttl).Unix()
	ip := clientIP(r)
	payload := fmt.Sprintf("%s|%d|%s|%s", cookieVersion, exp, hex.EncodeToString(nonce), ip)
	mac := hmac.New(sha256.New, c.secret)
	mac.Write([]byte(payload))
	sig := hex.EncodeToString(mac.Sum(nil))
	value := base64.RawURLEncoding.EncodeToString([]byte(payload + "|" + sig))

	http.SetCookie(w, &http.Cookie{
		Name:     c.name,
		Value:    value,
		Path:     "/",
		HttpOnly: true,
		Secure:   c.secure || r.TLS != nil || strings.EqualFold(r.Header.Get("X-Forwarded-Proto"), "https"),
		SameSite: http.SameSiteLaxMode,
		Expires:  time.Unix(exp, 0),
		MaxAge:   int(c.ttl.Seconds()),
	})
	return nil
}

func (c *cookieSigner) valid(r *http.Request) bool {
	cookie, err := r.Cookie(c.name)
	if err != nil || cookie.Value == "" {
		return false
	}
	raw, err := base64.RawURLEncoding.DecodeString(cookie.Value)
	if err != nil {
		return false
	}
	parts := strings.Split(string(raw), "|")
	if len(parts) != 5 {
		return false
	}
	ver, expStr, nonce, ip, sig := parts[0], parts[1], parts[2], parts[3], parts[4]
	if ver != cookieVersion || nonce == "" || sig == "" {
		return false
	}
	exp, err := strconv.ParseInt(expStr, 10, 64)
	if err != nil || time.Now().Unix() > exp {
		return false
	}
	if ip != clientIP(r) {
		return false
	}
	payload := fmt.Sprintf("%s|%s|%s|%s", ver, expStr, nonce, ip)
	mac := hmac.New(sha256.New, c.secret)
	mac.Write([]byte(payload))
	expected := hex.EncodeToString(mac.Sum(nil))
	return hmac.Equal([]byte(expected), []byte(sig))
}

func clientIP(r *http.Request) string {
	if ip, ok := caddyhttp.GetVar(r.Context(), caddyhttp.ClientIPVarKey).(string); ok && ip != "" {
		return ip
	}
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}