package anubis

import (
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/sync/singleflight"
)

const (
	// AuthCookie is the signed-JWT cookie Anubis sets once a challenge passes.
	AuthCookie = "techaro.lol-anubis-auth"
	// passChallengePath is the endpoint that validates a solved challenge.
	passChallengePath = "/.within.website/x/cmd/anubis/api/pass-challenge"
	// DefaultUserAgent is a stable, browser-like UA. The Anubis JWT is bound to
	// request metadata (incl. User-Agent), so every request through this
	// transport must use the same value.
	DefaultUserAgent = "Mozilla/5.0 (X11; Linux x86_64; rv:128.0) Gecko/20100101 Firefox/128.0"
)

// maxBodyPeek caps how much of a response body we read to detect/parse a
// challenge. Anubis interstitials are small; real pages can be large, so we
// only buffer when the response looks like an interstitial.
const maxChallengePeek = 1 << 20 // 1 MiB

// Transport is an http.RoundTripper that transparently solves Anubis
// proof-of-work challenges. When a request hits an interstitial, it solves the
// PoW, calls pass-challenge to obtain the auth cookie (stored in a shared
// jar), then replays the original request.
//
// A singleflight group collapses concurrent solves so only one goroutine pays
// the PoW cost while others wait for the resulting cookie.
type Transport struct {
	Base      http.RoundTripper
	UserAgent string
	Jar       http.CookieJar
	Logger    *slog.Logger

	sf singleflight.Group
	mu sync.Mutex // guards client construction
	hc *http.Client
}

// New returns a Transport with a fresh cookie jar and default settings.
func New() (*Transport, error) {
	jar, err := cookiejar.New(nil)
	if err != nil {
		return nil, fmt.Errorf("anubis: create cookie jar: %w", err)
	}
	return &Transport{
		Base:      http.DefaultTransport,
		UserAgent: DefaultUserAgent,
		Jar:       jar,
		Logger:    slog.Default(),
	}, nil
}

// Client returns an *http.Client that uses this Transport and shares its cookie
// jar, so callers benefit from both the auto-solving and the persisted cookie.
func (t *Transport) Client() *http.Client {
	t.mu.Lock()
	defer t.mu.Unlock()
	if t.hc == nil {
		t.hc = &http.Client{Transport: t, Jar: t.Jar, Timeout: 30 * time.Second}
	}
	return t.hc
}

// RoundTrip implements http.RoundTripper.
func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
	t.setUA(req)

	resp, err := t.base().RoundTrip(req)
	if err != nil {
		return nil, err
	}

	// Anubis serves interstitials as 200/text-html. Cheap content-type gate
	// before we buffer anything.
	if !maybeHTML(resp) {
		return resp, nil
	}

	body, err := io.ReadAll(io.LimitReader(resp.Body, maxChallengePeek))
	if err != nil {
		resp.Body.Close()
		return nil, fmt.Errorf("anubis: read body: %w", err)
	}
	resp.Body.Close()

	if !isInterstitial(body) {
		// Not a challenge; hand back a response with a replayable body.
		return withBody(resp, body), nil
	}

	if t.Logger != nil {
		t.Logger.Debug("anubis interstitial detected", "url", req.URL.String())
	}
	// The interstitial sets a cookie-verification cookie that pass-challenge
	// requires. The outer http.Client only commits response cookies to the jar
	// after RoundTrip returns, so capture them now — before we solve.
	t.Jar.SetCookies(req.URL, resp.Cookies())

	if err := t.solveFor(req, body); err != nil {
		return nil, err
	}

	// Replay the original request now that we hold the auth cookie. The replay
	// goes through the base transport directly, so attach jar cookies manually.
	replay := req.Clone(req.Context())
	t.setUA(replay)
	t.applyCookies(replay)
	resp2, err := t.base().RoundTrip(replay)
	if err != nil {
		return nil, err
	}
	if maybeHTML(resp2) {
		body2, err := io.ReadAll(io.LimitReader(resp2.Body, maxChallengePeek))
		resp2.Body.Close()
		if err != nil {
			return nil, fmt.Errorf("anubis: read replayed body: %w", err)
		}
		if isInterstitial(body2) {
			return nil, fmt.Errorf("anubis: still challenged after passing challenge for %s", req.URL)
		}
		return withBody(resp2, body2), nil
	}
	return resp2, nil
}

// solveFor parses the challenge from body, computes the answer, and calls
// pass-challenge to install the auth cookie. Concurrent calls for the same
// host are collapsed via singleflight.
func (t *Transport) solveFor(req *http.Request, body []byte) error {
	key := req.URL.Host
	_, err, _ := t.sf.Do(key, func() (any, error) {
		c, err := parseChallenge(body)
		if err != nil {
			return nil, err
		}
		passURL, err := t.answer(req, c)
		if err != nil {
			return nil, err
		}
		return nil, t.passChallenge(req, passURL)
	})
	return err
}

// answer computes the pass-challenge URL for the given challenge, performing
// any required wait or proof-of-work. It supports the "preact" challenge (a
// plain SHA-256 of randomData plus a timing gate) and the "fast"/"slow"
// proof-of-work challenges (nonce search).
func (t *Transport) answer(orig *http.Request, c *challenge) (string, error) {
	start := time.Now()
	switch c.algorithm {
	case "preact":
		// Answer is SHA-256(randomData); the server enforces a minimum wait of
		// difficulty*80ms. Mirror the frontend's difficulty*125ms delay.
		result := sha256Hex(c.randomData)
		wait := time.Duration(c.difficulty) * 125 * time.Millisecond
		time.Sleep(wait)

		u := resolveRef(orig.URL, c.redir)
		q := u.Query()
		q.Set("result", result)
		u.RawQuery = q.Encode()
		if t.Logger != nil {
			t.Logger.Info("answered anubis preact challenge",
				"difficulty", c.difficulty, "waited", wait)
		}
		return u.String(), nil

	case "fast", "slow", "proofofwork", "":
		hash, nonce := Solve(c.randomData, c.difficulty)
		u := &url.URL{Scheme: orig.URL.Scheme, Host: orig.URL.Host, Path: passChallengePath}
		q := url.Values{
			"id":          {c.id},
			"response":    {hash},
			"nonce":       {strconv.FormatUint(nonce, 10)},
			"redir":       {orig.URL.RequestURI()},
			"elapsedTime": {strconv.FormatInt(time.Since(start).Milliseconds(), 10)},
		}
		u.RawQuery = q.Encode()
		if t.Logger != nil {
			t.Logger.Info("solved anubis pow challenge",
				"difficulty", c.difficulty, "nonce", nonce, "elapsed", time.Since(start))
		}
		return u.String(), nil

	default:
		return "", fmt.Errorf("anubis: unsupported challenge algorithm %q", c.algorithm)
	}
}

// passChallenge GETs the pass-challenge URL so Anubis sets the auth cookie in
// the shared jar.
func (t *Transport) passChallenge(orig *http.Request, passURL string) error {
	preq, err := http.NewRequestWithContext(orig.Context(), http.MethodGet, passURL, nil)
	if err != nil {
		return fmt.Errorf("anubis: build pass-challenge request: %w", err)
	}
	t.setUA(preq)
	preq.Header.Set("Referer", orig.URL.String())

	// Don't follow the post-pass redirect; we only need the Set-Cookie. Use a
	// client bound to the shared jar so the cookie is captured.
	hc := &http.Client{
		Transport: t.base(),
		Jar:       t.Jar,
		CheckRedirect: func(*http.Request, []*http.Request) error {
			return http.ErrUseLastResponse
		},
		Timeout: 30 * time.Second,
	}
	resp, err := hc.Do(preq)
	if err != nil {
		return fmt.Errorf("anubis: pass-challenge: %w", err)
	}
	defer resp.Body.Close()
	io.Copy(io.Discard, io.LimitReader(resp.Body, maxChallengePeek))

	if !t.hasAuthCookie(orig.URL) {
		return fmt.Errorf("anubis: pass-challenge did not yield %s cookie (status %d)", AuthCookie, resp.StatusCode)
	}
	return nil
}

// hasAuthCookie reports whether the jar holds an Anubis auth cookie for u.
// Deployments customize the cookie prefix (default "techaro.lol-anubis-auth",
// root-me uses "anubis-cookie-auth"), so we match any non-empty cookie whose
// name ends in "-auth" while excluding the "-cookie-verification" helper.
func (t *Transport) hasAuthCookie(u *url.URL) bool {
	for _, c := range t.Jar.Cookies(u) {
		if c.Value == "" || strings.Contains(c.Name, "verification") {
			continue
		}
		if c.Name == AuthCookie || strings.HasSuffix(c.Name, "-auth") {
			return true
		}
	}
	return false
}

func (t *Transport) base() http.RoundTripper {
	if t.Base != nil {
		return t.Base
	}
	return http.DefaultTransport
}

func (t *Transport) setUA(req *http.Request) {
	ua := t.UserAgent
	if ua == "" {
		ua = DefaultUserAgent
	}
	req.Header.Set("User-Agent", ua)
	if req.Header.Get("Accept-Language") == "" {
		req.Header.Set("Accept-Language", "en-US,en;q=0.9")
	}
}