dsfree2api/internal/upstream/client.go

package upstream

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"html"
	"io"
	"log/slog"
	"net/http"
	"os"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	"github.com/nyoungo/dsfree2api/internal/config"
	"github.com/nyoungo/dsfree2api/internal/httpx"
	"github.com/nyoungo/dsfree2api/internal/turnstile"
)

type Route struct {
	Name  string
	Proxy string // "" means direct
}

func (r Route) Key() string {
	if r.Proxy == "" {
		return "direct"
	}
	return r.Proxy
}

// RouteCandidate is one egress route supplied by the proxy pool.
type RouteCandidate struct {
	Name  string
	Proxy string
}

// RouteSource resolves per-site route candidates and records request
// outcomes. Implemented by proxypool.Manager; nil when the pool is unused.
type RouteSource interface {
	Candidates(site string) []RouteCandidate
	Report(proxy string, ok bool)
}

type chatConfig struct {
	BotID     int
	PostID    int
	Nonce     string
	AJAXURL   string
	FetchedAt time.Time
}

type Client struct {
	cfg  *config.Config
	ts   *turnstile.Solver
	log  *slog.Logger
	pool RouteSource

	stateMu  sync.Mutex
	chatCfg  map[string]chatConfig
	cfgLocks map[string]*sync.Mutex
	gates    map[string]*siteGate

	// quotaCool marks sites that recently answered with quota errors; they
	// are tried after the healthy mirrors until the cooldown expires.
	quotaCool map[string]time.Time
	// lastWarn* dedupes repeated identical upstream errors so a quota
	// storm cannot drown the log.
	lastWarnMsg string
	lastWarnAt  time.Time

	// quotaObserver, when set, is notified on the first site-side quota
	// exhaustion of each request — including ones hidden by cross-site
	// failover.
	quotaObserver func(site, model, route string, err error)

	// chatOnceOverride is a test seam; nil in production.
	chatOnceOverride func(ctx context.Context, site config.Site, modelID string, model config.Model,
		prompt string, route Route, info *ServeInfo, yield func(Event) error) error
}

func New(cfg *config.Config, ts *turnstile.Solver, log *slog.Logger, pool RouteSource) *Client {
	if log == nil {
		log = slog.Default()
	}
	return &Client{
		cfg:       cfg,
		ts:        ts,
		log:       log,
		pool:      pool,
		chatCfg:   map[string]chatConfig{},
		cfgLocks:  map[string]*sync.Mutex{},
		gates:     map[string]*siteGate{},
		quotaCool: map[string]time.Time{},
	}
}

// SetQuotaObserver registers a callback that fires whenever a site reports
// quota exhaustion, even when cross-site failover later succeeds and the error
// would otherwise never surface. Set it before serving requests.
func (c *Client) SetQuotaObserver(fn func(site, model, route string, err error)) {
	c.quotaObserver = fn
}

// Routes returns the primary proxy followed by the fallbacks, deduplicated.
func (c *Client) Routes() []Route {
	c.cfg.RLock()
	primary := strings.TrimSpace(c.cfg.Proxy.URL)
	fallbacks := append([]string(nil), c.cfg.Proxy.FallbackURLs...)
	c.cfg.RUnlock()

	type cand struct {
		name  string
		proxy string
	}
	list := []cand{{"primary", primary}}
	for i, u := range fallbacks {
		list = append(list, cand{fmt.Sprintf("fallback-%d", i+1), strings.TrimSpace(u)})
	}
	seen := map[string]bool{}
	out := make([]Route, 0, len(list))
	for _, it := range list {
		key := it.proxy
		if key == "" {
			key = "direct"
		}
		if seen[key] {
			continue
		}
		seen[key] = true
		out = append(out, Route{Name: it.name, Proxy: it.proxy})
	}
	return out
}

// quotaCoolTTL is how long a site that answered with quota errors keeps
// being tried after its healthy mirrors.
const quotaCoolTTL = 10 * time.Minute

// quotaCooled reports whether a site is inside its quota cooldown window.
func (c *Client) quotaCooled(code string) bool {
	c.stateMu.Lock()
	defer c.stateMu.Unlock()
	until, ok := c.quotaCool[code]
	return ok && time.Now().Before(until)
}

// setQuotaCool cools a site down after it reported quota exhaustion.
// Repeat detections inside the window only extend it; the INFO line marks
// the transition, not every probe that confirms it.
func (c *Client) setQuotaCool(code string) {
	c.stateMu.Lock()
	_, wasCooled := c.quotaCool[code]
	// 每次都把截止时间顺延，重复探测只延长窗口；只有从"未冷却"进入
	// "冷却"这一次状态翻转才记 INFO，之后的确认探测不再刷日志。
	c.quotaCool[code] = time.Now().Add(quotaCoolTTL)
	c.stateMu.Unlock()
	if !wasCooled {
		c.log.Info("site quota cooldown", "site", code, "for", quotaCoolTTL)
	}
}

// sortQuotaCoolLast moves quota-cooled sites behind the healthy candidates,
// preserving relative order inside both groups.
func (c *Client) sortQuotaCoolLast(ids []string) []string {
	var cooled, rest []string
	for _, id := range ids {
		m, ok := c.modelCopy(id)
		if ok && c.quotaCooled(m.Site) {
			cooled = append(cooled, id)
			continue
		}
		rest = append(rest, id)
	}
	if len(cooled) == 0 {
		return ids
	}
	return append(rest, cooled...)
}

// clearQuotaCool lifts a site's cooldown once it serves again after a quota
// error — the fresh visitor identity restored its daily tier.
func (c *Client) clearQuotaCool(code string) {
	c.stateMu.Lock()
	_, ok := c.quotaCool[code]
	if ok {
		delete(c.quotaCool, code)
	}
	c.stateMu.Unlock()
	if ok {
		c.log.Info("site quota restored", "site", code)
	}
}

// probeBalanceAsync logs the site's current free tier after a quota
// recovery. Best effort — failures only show up at DEBUG.
func (c *Client) probeBalanceAsync(site config.Site, model config.Model, route Route) {
	if c.ts == nil {
		return
	}
	gid, ok := c.ts.GuestID(site.Code, route.Proxy)
	if !ok {
		return
	}
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
		defer cancel()
		bal, err := c.FetchBalance(ctx, site, gid, model.BotID, route)
		if err != nil {
			c.log.Debug("balance probe failed", "site", site.Code, "error", err)
			return
		}
		c.log.Info("site balance", "site", site.Code,
			"remaining", bal.Free.Remaining, "limit", bal.Free.Limit)
	}()
}

// logUpstreamErr records one failed upstream attempt. Quota errors cool the
// site down and repeats of an identical error inside a short window drop to
// DEBUG, so a balance-starved site neither taxes every continuation round
// nor floods the log.
func (c *Client) logUpstreamErr(site, modelID string, attempt int, route Route, err error) {
	if isQuota(err) {
		c.setQuotaCool(site)
	}
	msg := err.Error()
	c.stateMu.Lock()
	dup := msg == c.lastWarnMsg && time.Since(c.lastWarnAt) < 30*time.Second
	c.lastWarnMsg = msg
	c.lastWarnAt = time.Now()
	c.stateMu.Unlock()
	if dup {
		c.log.Debug("upstream error (repeating)", "model", modelID, "attempt", attempt+1, "route", route.Name, "error", err)
		return
	}
	c.log.Warn("upstream error", "model", modelID, "attempt", attempt+1, "route", route.Name, "error", err)
}

// logCacheReject fingerprints a prompt the site's cache step called empty:
// byte class counts plus head/tail quotes tell whether the rejection tracks
// prompt content (WAF/filter) or session state.
func (c *Client) logCacheReject(site, modelID, prompt string) {
	nul, ctrl := 0, 0
	for i := 0; i < len(prompt); i++ {
		b := prompt[i]
		switch {
		case b == 0:
			nul++
		case b < 0x20 && b != '\n' && b != '\r' && b != '\t':
			ctrl++
		}
	}
	head, tail := prompt, ""
	if r := []rune(prompt); len(r) > 160 {
		head, tail = string(r[:80]), string(r[len(r)-80:])
	}
	sum := sha256.Sum256([]byte(prompt))
	c.log.Warn("cache message rejected",
		"site", site, "model", modelID, "chars", len(prompt),
		"nul", nul, "ctrl", ctrl, "sha", hex.EncodeToString(sum[:8]),
		"head", strconv.Quote(head), "tail", strconv.Quote(tail))
}

// routesFor returns the ordered candidate routes for one site: proxy-pool
// candidates first (sticky, health filtered), then the global primary and
// fallback lines as the safety net.
func (c *Client) routesFor(siteCode string) []Route {
	out := make([]Route, 0, 4)
	seen := map[string]bool{}
	add := func(r Route) {
		key := r.Key()
		if seen[key] {
			return
		}
		seen[key] = true
		out = append(out, r)
	}
	if c.pool != nil {
		for _, cand := range c.pool.Candidates(siteCode) {
			add(Route{Name: cand.Name, Proxy: cand.Proxy})
		}
	}
	for _, r := range c.Routes() {
		add(r)
	}
	return out
}

// poolReport feeds a route outcome back to the proxy pool; it is a no-op for
// global routes and the direct connection.
func (c *Client) poolReport(proxy string, ok bool) {
	if c.pool == nil || proxy == "" {
		return
	}
	c.pool.Report(proxy, ok)
}

// Chat streams the prompt through the upstream site, honouring route
// failover, slow-start detection, session refresh and cross-site failover.
func (c *Client) Chat(ctx context.Context, modelID string, prompt string, info *ServeInfo, yield func(Event) error) error {
	model, ok := c.modelCopy(modelID)
	if !ok {
		return errf("unknown model %s", modelID)
	}
	if !model.Enabled {
		return errf("model %s is disabled", modelID)
	}

	candidates := []string{modelID}
	if c.cfg.Upstream.CrossSiteFailover {
		for _, alt := range c.alternates(modelID, model) {
			candidates = append(candidates, alt)
		}
	}
	// A site that recently reported quota errors keeps being tried, but
	// behind the healthy mirrors — continuation rounds stop paying its
	// refresh tax while its balance is empty.
	candidates = c.sortQuotaCoolLast(candidates)

	var last error
	triedPrimary := false
	prev := ""
	for _, id := range candidates {
		m := model
		if id == modelID {
			triedPrimary = true
		} else {
			mc, ok := c.modelCopy(id)
			if !ok || !mc.Enabled {
				continue
			}
			m = mc
		}
		// Only a real switch after a failed candidate is a failover; the
		// first pick of a request (even a sorted-past primary) is not.
		if prev != "" && prev != id {
			c.log.Warn("cross-site failover", "from", prev, "to", id)
		}
		err := c.chatWithRetries(ctx, id, m, prompt, info, yield)
		if err == nil {
			return nil
		}
		last = err
		if contentStarted(err) || !turnstile.Retryable(err) {
			return AsUpstream(err)
		}
		// Quota means this site has no balance and a cache-rejection means
		// its session is off: the next candidate is the point. Other errors
		// keep going only while the primary has not had its chance yet (a
		// quota-cooled primary sits at the end of the list); once it failed
		// too, report that error as before.
		if !isQuota(err) && !isCacheEmptyProblem(err) && triedPrimary {
			return AsUpstream(err)
		}
		prev = id
	}
	return AsUpstream(last)
}

// alternates lists other enabled sites hosting the same upstream model.
func (c *Client) alternates(modelID string, m config.Model) []string {
	c.cfg.RLock()
	defer c.cfg.RUnlock()
	var out []string
	for id, other := range c.cfg.Models {
		if id == modelID || other.Site == m.Site {
			continue
		}
		if other.UpstreamID != m.UpstreamID || !other.Enabled {
			continue
		}
		if s, ok := c.cfg.Sites[other.Site]; !ok || !s.Enabled {
			continue
		}
		out = append(out, id)
	}
	return out
}

func (c *Client) modelCopy(id string) (config.Model, bool) {
	c.cfg.RLock()
	defer c.cfg.RUnlock()
	m, ok := c.cfg.Models[id]
	if !ok {
		return config.Model{}, false
	}
	return *m, true
}

func (c *Client) siteCopy(code string) (config.Site, bool) {
	c.cfg.RLock()
	defer c.cfg.RUnlock()
	s, ok := c.cfg.Sites[code]
	if !ok {
		return config.Site{}, false
	}
	return *s, true
}

func (c *Client) userAgent() string {
	c.cfg.RLock()
	defer c.cfg.RUnlock()
	return c.cfg.Upstream.UserAgent
}

func (c *Client) chatWithRetries(ctx context.Context, modelID string, model config.Model, prompt string, info *ServeInfo, yield func(Event) error) error {
	site, ok := c.siteCopy(model.Site)
	if !ok {
		return errf("unknown site %s", model.Site)
	}
	if !site.Enabled {
		return errf("site %s is disabled", model.Site)
	}
	routes := c.routesFor(site.Code)
	quotaNotified := false

	var lastErr error
	sawQuota := false
	for attempt := 0; ; attempt++ {
		var failed []Route
		for i, route := range routes {
			// A first-time Turnstile solve can far exceed the slow-start
			// budget; only arm the timer when cookies are already cached
			// (or no solver is configured at all).
			cookiesReady := c.ts == nil || c.ts.CookiesReady(site.Code, route.Proxy)
			allowSlow := i == 0 && len(routes) > 1 && cookiesReady
			err := c.chatOnceSlowStart(ctx, site, modelID, model, prompt, route, allowSlow, info, yield)
			if err == nil {
				if sawQuota {
					c.clearQuotaCool(site.Code)
					c.probeBalanceAsync(site, model, route)
				}
				c.poolReport(route.Proxy, true)
				return nil
			}
			lastErr = err
			if isQuota(err) && !quotaNotified && c.quotaObserver != nil {
				quotaNotified = true
				c.quotaObserver(site.Code, modelID, route.Name, err)
			}
			if contentStarted(err) {
				return AsUpstream(err)
			}
			// A solver-side rejection (bad key, rejected challenge) will fail
			// identically on every attempt — stop instead of burning retries.
			if !turnstile.Retryable(err) {
				return AsUpstream(err)
			}
			// The site demands Turnstile verification while the solver is
			// switched off: refreshing cookies cannot help, say why instead.
			if isSession(err) && (c.ts == nil || !c.ts.Enabled()) {
				return &UpstreamError{
					Msg: "site requires Turnstile verification — import cookies on the console " +
						"Turnstile page, or enable [turnstile] to solve automatically",
					Cause: err,
				}
			}
			if isConfigProblem(err) {
				return AsUpstream(err)
			}
			// The cache step rejected the message in a way config drops and
			// cookie refreshes do not fix; stop burning refresh cycles so
			// the caller can fail over to a mirror right away.
			if isCacheEmptyProblem(err) {
				return AsUpstream(err)
			}
			// Quota/session failures are site-side, not proxy-side — do not
			// blame the egress endpoint for them.
			if !isQuota(err) && !isSession(err) {
				c.poolReport(route.Proxy, false)
			}
			failed = append(failed, route)
			c.logUpstreamErr(site.Code, modelID, attempt, route, err)
			if isQuota(err) {
				sawQuota = true
			}
			if i < len(routes)-1 {
				continue
			}
			break
		}

		c.cfg.RLock()
		autoRefresh := c.cfg.Upstream.AutoRefresh
		backoff := c.cfg.Upstream.RetryBackoffSeconds
		retries := c.cfg.Upstream.RefreshRetries
		c.cfg.RUnlock()
		if !autoRefresh || attempt >= retries {
			break
		}
		if len(failed) > 0 {
			wait := time.Duration(float64(attempt+1) * backoff * float64(time.Second))
			select {
			case <-ctx.Done():
				return AsUpstream(ctx.Err())
			case <-time.After(wait):
			}
			// Session errors need a fresh cookie right away. Quota errors
			// often clear once the cached chat config is dropped, so the
			// first cycle retries cheaply and only the second one pays for
			// a Turnstile re-solve.
			forceCookie := isSession(lastErr) || (isQuota(lastErr) && attempt >= 1)
			// The free daily tier is keyed to the visitor cookie: swapping
			// in a brand-new visitor id restores it instantly, the same
			// effect as a fresh private window. Only a re-solve (attempt
			// >= 1) then remains as the slower fallback.
			if isQuota(lastErr) && attempt == 0 && c.ts != nil {
				if gid := c.ts.RotateGuest(site.Code); gid != "" {
					c.log.Info("rotated visitor identity", "site", site.Code, "reason", "quota")
				}
			}
			for _, route := range failed {
				if err := c.refreshSession(ctx, site, modelID, route, forceCookie); err != nil {
					c.log.Warn("session refresh failed", "model", modelID, "error", err)
				}
			}
		}
	}
	return AsUpstream(lastErr)
}

// chatOnceSlowStart bounds the time the primary route may take to produce its
// first upstream event before a fallback route is tried instead. Errors that
// happen after the first event are marked ContentStarted so callers never
// restart a response the client already partially received.
func (c *Client) chatOnceSlowStart(
	ctx context.Context,
	site config.Site,
	modelID string,
	model config.Model,
	prompt string,
	route Route,
	allowSlow bool,
	info *ServeInfo,
	yield func(Event) error,
) error {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	events := make(chan Event, 16)
	errCh := make(chan error, 1)
	gotFirst := false

	go func() {
		fn := c.chatOnce
		if c.chatOnceOverride != nil {
			fn = c.chatOnceOverride
		}
		err := fn(ctx, site, modelID, model, prompt, route, info, func(ev Event) error {
			select {
			case events <- ev:
				return nil
			case <-ctx.Done():
				return ctx.Err()
			}
		})
		close(events)
		errCh <- err
	}()

	mark := func(err error) error {
		if err == nil || !gotFirst {
			return err
		}
		ue := AsUpstream(err)
		ue.ContentStarted = true
		return ue
	}
	take := func(ev Event) error {
		gotFirst = true
		if err := yield(ev); err != nil {
			return mark(err)
		}
		return nil
	}

	if allowSlow {
		var slow float64
		c.cfg.RLock()
		slow = c.cfg.Proxy.SlowStartSeconds
		c.cfg.RUnlock()
		if slow > 0 {
			timer := time.NewTimer(time.Duration(slow * float64(time.Second)))
			defer timer.Stop()
			select {
			case ev, ok := <-events:
				if !ok {
					return mark(<-errCh)
				}
				if err := take(ev); err != nil {
					return err
				}
			case err := <-errCh:
				for ev := range events {
					if e := take(ev); e != nil {
						return e
					}
				}
				return mark(err)
			case <-timer.C:
				cancel()
				<-errCh
				return errf("primary route produced no upstream event within %gs", slow)
			}
		}
	}

	for {
		ev, ok := <-events
		if !ok {
			return mark(<-errCh)
		}
		if err := take(ev); err != nil {
			return err
		}
	}
}

func (c *Client) chatOnce(
	ctx context.Context,
	site config.Site,
	modelID string,
	model config.Model,
	prompt string,
	route Route,
	info *ServeInfo,
	yield func(Event) error,
) error {
	if err := c.acquire(ctx, site.Code); err != nil {
		return AsUpstream(err)
	}
	defer c.release(site.Code)

	if info != nil {
		info.Set(modelID, site.Code, route.Name)
	}

	c.cfg.RLock()
	streamTimeout := time.Duration(c.cfg.Upstream.StreamTimeout * float64(time.Second))
	timeout := time.Duration(c.cfg.Upstream.Timeout * float64(time.Second))
	ua := c.cfg.Upstream.UserAgent
	c.cfg.RUnlock()

	sess, err := httpx.NewSession(httpx.Options{ProxyURL: route.Proxy, Timeout: streamTimeout})
	if err != nil {
		return errf("create session: %v", err)
	}
	defer sess.Close()

	if err := c.ts.ApplyValidCookies(ctx, sess, &site, route.Proxy); err != nil {
		return AsUpstream(err)
	}

	cc, err := c.chatConfig(ctx, sess, site, modelID, model, route, ua)
	if err != nil {
		return err
	}

	referer := site.BaseURL + model.PagePath
	sessionID := newUUID()
	convUUID := newUUID()
	clientMsgID := fmt.Sprintf("aipkit-client-msg-%d-%d-%s", cc.BotID, time.Now().UnixMilli(), randHex(3))

	resp, err := sess.Do(ctx, httpx.Request{
		Method: "POST",
		URL:    cc.AJAXURL,
		Header: ajaxHeaders(site, referer, ua),
		Body: []byte(encodeForm(map[string]string{
			"action":                 "aipkit_cache_sse_message",
			"message":                prompt,
			"_ajax_nonce":            cc.Nonce,
			"bot_id":                 itoa(cc.BotID),
			"session_id":             sessionID,
			"conversation_uuid":      convUUID,
			"user_client_message_id": clientMsgID,
		})),
		Timeout: timeout,
	})
	if err != nil {
		return errf("cache message request failed: %v", err)
	}
	if bytes.Contains(resp.Body, []byte("empty_data_to_cache")) {
		c.logCacheReject(site.Code, modelID, prompt)
	}
	if resp.Status >= 400 {
		if resp.Status == 403 || resp.Status == 401 {
			return newSession(fmt.Sprintf("cache message http %d", resp.Status))
		}
		return errf("cache message http %d: %s", resp.Status, truncate(string(resp.Body), 400))
	}

	var init struct {
		Success bool            `json:"success"`
		Data    json.RawMessage `json:"data"`
	}
	if err := json.Unmarshal(resp.Body, &init); err != nil {
		return errf("cache message bad json: %v", err)
	}
	if !init.Success {
		if code := extractCode(init.Data); code != "" && strings.Contains(strings.ToLower(code), "nonce") {
			c.dropChatCfg(modelID, route)
		}
		payload := map[string]any{}
		_ = json.Unmarshal(resp.Body, &payload)
		if isSessionPayload(payload) {
			return newSession("cache message session failed: " + truncate(string(resp.Body), 400))
		}
		return errf("cache message failed: %s", truncate(string(resp.Body), 400))
	}
	var dataObj struct {
		CacheKey string `json:"cache_key"`
	}
	if err := json.Unmarshal(init.Data, &dataObj); err != nil || dataObj.CacheKey == "" {
		return errf("cache message missing cache_key: %s", truncate(string(resp.Body), 400))
	}

	params := encodeForm(map[string]string{
		"action":            "aipkit_frontend_chat_stream",
		"cache_key":         dataObj.CacheKey,
		"bot_id":            itoa(cc.BotID),
		"session_id":        sessionID,
		"conversation_uuid": convUUID,
		"post_id":           itoa(cc.PostID),
		"_ts":               itoa64(time.Now().UnixMilli()),
		"_ajax_nonce":       cc.Nonce,
	})

	streamCtx, streamCancel := context.WithTimeout(ctx, streamTimeout)
	defer streamCancel()

	headers := ajaxHeaders(site, referer, ua)
	headers["Accept"] = "text/event-stream"

	rc, status, err := sess.Stream(streamCtx, httpx.Request{
		Method: "GET",
		URL:    cc.AJAXURL + "?" + params,
		Header: headers,
	})
	if err != nil {
		if streamCtx.Err() != nil && ctx.Err() == nil {
			return errf("upstream stream timeout: %v", err)
		}
		return streamHTTPError(status, err)
	}
	defer rc.Close()

	if os.Getenv("DSFREE_SSE_TIMING") != "" {
		start := time.Now()
		rc = &sseTimingReader{rc: rc, start: start, logf: func(n int) {
			c.log.Info("sse read", "at_ms", time.Since(start).Milliseconds(), "n", n)
		}}
	}

	return readSSE(rc, func(event string, lines []string) error {
		events, err := translateEvent(event, lines)
		if err != nil {
			return err
		}
		for _, ev := range events {
			if err := yield(ev); err != nil {
				return err
			}
		}
		return nil
	})
}

// streamHTTPError classifies a stream setup failure by its HTTP status. 401/403
// mean the cached session is no longer accepted (expired nonce / Turnstile
// cookie), which must be reported as a session error so the caller refreshes
// cookies instead of blaming the egress endpoint (poolReport(false)) and
// keeping the stale cookie.
func streamHTTPError(status int, err error) error {
	if status == http.StatusUnauthorized || status == http.StatusForbidden {
		return newSession(fmt.Sprintf("upstream stream http %d: %v", status, err))
	}
	return AsUpstream(err)
}

// sseTimingReader (diagnosis only, gated by DSFREE_SSE_TIMING) logs when bytes
// arrive from the upstream stream.
type sseTimingReader struct {
	rc    io.ReadCloser
	start time.Time
	logf  func(n int)
}

func (r *sseTimingReader) Read(p []byte) (int, error) {
	n, err := r.rc.Read(p)
	if n > 0 {
		r.logf(n)
	}
	return n, err
}

func (r *sseTimingReader) Close() error { return r.rc.Close() }

// chatConfig fetches (and caches) the page-embedded nonce for a model.
func (c *Client) chatConfig(
	ctx context.Context,
	sess httpx.Session,
	site config.Site,
	modelID string,
	model config.Model,
	route Route,
	ua string,
) (chatConfig, error) {
	key := modelID + "|" + route.Key()
	if v, ok := c.loadChatCfg(key); ok {
		return v, nil
	}
	lock := c.lockFor(key)
	lock.Lock()
	defer lock.Unlock()
	if v, ok := c.loadChatCfg(key); ok {
		return v, nil
	}

	pageURL := site.BaseURL + model.PagePath
	resp, err := sess.Do(ctx, httpx.Request{
		Method:  "GET",
		URL:     pageURL,
		Header:  browserHeaders(site, ua),
		Timeout: 60 * time.Second,
	})
	if err != nil {
		return chatConfig{}, errf("chat config fetch failed: %v", err)
	}
	if resp.Status >= 400 {
		return chatConfig{}, errf("chat config http %d", resp.Status)
	}

	parsed, found := parseChatConfig(string(resp.Body), site, model)
	if !found {
		if model.BotID <= 0 || model.PostID <= 0 {
			return chatConfig{}, errf("no chat data-config found for %s", modelID)
		}
		nonce, err := c.fetchNonce(ctx, sess, site, model, ua)
		if err != nil {
			return chatConfig{}, err
		}
		parsed = chatConfig{
			BotID: model.BotID, PostID: model.PostID, Nonce: nonce,
			AJAXURL: site.AJAXURL, FetchedAt: time.Now(),
		}
	}
	parsed.FetchedAt = time.Now()
	c.storeChatCfg(key, parsed)
	c.log.Info("loaded chat config", "model", modelID, "route", route.Name,
		"bot_id", parsed.BotID, "post_id", parsed.PostID)
	return parsed, nil
}

var containerRe = regexp.MustCompile(`id="aipkit_chat_container_(\d+)"[^>]*data-config='([^']+)'`)

func parseChatConfig(page string, site config.Site, model config.Model) (chatConfig, bool) {
	for _, m := range containerRe.FindAllStringSubmatch(page, -1) {
		botID := atoi(m[1])
		raw := html.UnescapeString(m[2])
		var cfg struct {
			BotID   any    `json:"botId"`
			PostID  any    `json:"postId"`
			Nonce   string `json:"nonce"`
			AjaxURL string `json:"ajaxUrl"`
		}
		if err := json.Unmarshal([]byte(raw), &cfg); err != nil {
			continue
		}
		if model.BotID > 0 && botID != model.BotID {
			continue
		}
		out := chatConfig{
			BotID:     toInt(cfg.BotID, botID),
			PostID:    toInt(cfg.PostID, model.PostID),
			Nonce:     cfg.Nonce,
			AJAXURL:   cfg.AjaxURL,
			FetchedAt: time.Now(),
		}
		if out.AJAXURL == "" {
			out.AJAXURL = site.AJAXURL
		}
		if out.Nonce != "" {
			return out, true
		}
	}
	return chatConfig{}, false
}

func (c *Client) fetchNonce(ctx context.Context, sess httpx.Session, site config.Site, model config.Model, ua string) (string, error) {
	resp, err := sess.Do(ctx, httpx.Request{
		Method:  "POST",
		URL:     site.AJAXURL,
		Header:  ajaxHeaders(site, site.BaseURL+model.PagePath, ua),
		Body:    []byte(encodeForm(map[string]string{"action": "aipkit_get_frontend_chat_nonce", "bot_id": itoa(model.BotID)})),
		Timeout: 30 * time.Second,
	})
	if err != nil {
		return "", errf("nonce fetch failed: %v", err)
	}
	var out struct {
		Success bool `json:"success"`
		Data    struct {
			Nonce string `json:"nonce"`
		} `json:"data"`
	}
	if err := json.Unmarshal(resp.Body, &out); err != nil || !out.Success || out.Data.Nonce == "" {
		return "", errf("nonce fetch failed for %s: %s", model.ID, truncate(string(resp.Body), 400))
	}
	return out.Data.Nonce, nil
}

// refreshSession drops cached state for a route and optionally re-solves Turnstile.
func (c *Client) refreshSession(
	ctx context.Context,
	site config.Site,
	modelID string,
	route Route,
	invalidateCookies bool,
) error {
	c.dropChatCfg(modelID, route)
	if !invalidateCookies {
		return nil
	}
	c.log.Info("invalidate verified cookie cache", "site", site.Code, "route", route.Name)
	c.ts.Invalidate(site.Code)

	sess, err := httpx.NewSession(httpx.Options{ProxyURL: route.Proxy, Timeout: 60 * time.Second})
	if err != nil {
		return err
	}
	defer sess.Close()
	return c.ts.Refresh(ctx, sess, &site, route.Proxy)
}

// InvalidateAll clears cached nonces and cookies (web console action).
func (c *Client) InvalidateAll() {
	c.stateMu.Lock()
	c.chatCfg = map[string]chatConfig{}
	c.stateMu.Unlock()
	c.cfg.RLock()
	codes := make([]string, 0, len(c.cfg.Sites))
	for code := range c.cfg.Sites {
		codes = append(codes, code)
	}
	c.cfg.RUnlock()
	for _, code := range codes {
		c.ts.Invalidate(code)
	}
}

func (c *Client) ConfigCacheSize() int {
	c.stateMu.Lock()
	defer c.stateMu.Unlock()
	return len(c.chatCfg)
}

// ── cache helpers ────────────────────────────────────────────────

func (c *Client) loadChatCfg(key string) (chatConfig, bool) {
	c.stateMu.Lock()
	defer c.stateMu.Unlock()
	v, ok := c.chatCfg[key]
	if !ok {
		return chatConfig{}, false
	}
	c.cfg.RLock()
	ttl := time.Duration(c.cfg.Upstream.ConfigTTLSeconds) * time.Second
	c.cfg.RUnlock()
	if ttl <= 0 || time.Since(v.FetchedAt) > ttl {
		return chatConfig{}, false
	}
	return v, true
}

func (c *Client) storeChatCfg(key string, v chatConfig) {
	c.stateMu.Lock()
	defer c.stateMu.Unlock()
	c.chatCfg[key] = v
}

func (c *Client) dropChatCfg(modelID string, route Route) {
	suffix := "|" + route.Key()
	c.stateMu.Lock()
	defer c.stateMu.Unlock()
	for k := range c.chatCfg {
		if strings.HasSuffix(k, suffix) && strings.TrimSuffix(k, suffix) == modelID {
			delete(c.chatCfg, k)
		}
	}
}

func (c *Client) lockFor(key string) *sync.Mutex {
	c.stateMu.Lock()
	defer c.stateMu.Unlock()
	l, ok := c.cfgLocks[key]
	if !ok {
		l = &sync.Mutex{}
		c.cfgLocks[key] = l
	}
	return l
}

// ── concurrency gate ────────────────────────────────────────────

// siteGate 是单个站点的并发闸门。名额用 busy 计数而不是"首次使用时定容的
// channel"，这样 [limits].max_concurrent_per_site 被修改后立即生效，无需重启
// （channel 版本只在第一次 acquire 时定容，之后改配置一直不生效）。
type siteGate struct {
	mu   sync.Mutex
	busy int
	// wake 在每次 release 时被关闭并重建，用来广播唤醒等待者；等待者醒来后
	// 按当时最新的配置重新竞争名额。
	wake chan struct{}
}

func newSiteGate() *siteGate { return &siteGate{wake: make(chan struct{})} }

func (c *Client) acquire(ctx context.Context, site string) error {
	c.stateMu.Lock()
	g, ok := c.gates[site]
	if !ok {
		g = newSiteGate()
		c.gates[site] = g
	}
	c.stateMu.Unlock()
	return c.admit(ctx, site, g)
}

// admit 在当前配置的上限内占一个名额；超限时阻塞到有空位或 ctx 结束。
func (c *Client) admit(ctx context.Context, site string, g *siteGate) error {
	for {
		c.cfg.RLock()
		max := c.cfg.Limits.MaxConcurrentPerSite
		c.cfg.RUnlock()

		g.mu.Lock()
		if max <= 0 || g.busy < max {
			g.busy++
			g.mu.Unlock()
			return nil
		}
		wake := g.wake
		g.mu.Unlock()

		select {
		case <-ctx.Done():
			return fmt.Errorf("site %s busy: %w", site, ctx.Err())
		case <-wake:
		}
	}
}

func (c *Client) release(site string) {
	c.stateMu.Lock()
	g, ok := c.gates[site]
	c.stateMu.Unlock()
	if !ok {
		return
	}
	g.mu.Lock()
	if g.busy > 0 {
		g.busy--
	}
	close(g.wake)
	g.wake = make(chan struct{})
	g.mu.Unlock()
}

// ── helpers ─────────────────────────────────────────────────────

func browserHeaders(site config.Site, ua string) map[string]string {
	if ua == "" {
		ua = httpx.DefaultUserAgent
	}
	return map[string]string{
		"User-Agent":      ua,
		"Accept":          "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
		"Accept-Language": site.Language,
		"Origin":          site.BaseURL,
		"Referer":         site.BaseURL + "/",
	}
}

func ajaxHeaders(site config.Site, referer, ua string) map[string]string {
	if ua == "" {
		ua = httpx.DefaultUserAgent
	}
	return map[string]string{
		"User-Agent":       ua,
		"Accept":           "application/json, text/javascript, */*; q=0.01",
		"Accept-Language":  site.Language,
		"X-Requested-With": "XMLHttpRequest",
		"Origin":           site.BaseURL,
		"Referer":          referer,
	}
}

func encodeForm(v map[string]string) string {
	keys := make([]string, 0, len(v))
	for k := range v {
		keys = append(keys, k)
	}
	sortStrings(keys)
	var b strings.Builder
	for i, k := range keys {
		if i > 0 {
			b.WriteByte('&')
		}
		b.WriteString(urlEncode(k))
		b.WriteByte('=')
		b.WriteString(urlEncode(v[k]))
	}
	return b.String()
}

func urlEncode(s string) string {
	var b strings.Builder
	for i := 0; i < len(s); i++ {
		ch := s[i]
		switch {
		case ch >= 'a' && ch <= 'z', ch >= 'A' && ch <= 'Z', ch >= '0' && ch <= '9',
			ch == '-', ch == '_', ch == '.', ch == '~':
			b.WriteByte(ch)
		case ch == ' ':
			b.WriteByte('+')
		default:
			fmt.Fprintf(&b, "%%%02X", ch)
		}
	}
	return b.String()
}

func sortStrings(s []string) {
	for i := 1; i < len(s); i++ {
		for j := i; j > 0 && s[j] < s[j-1]; j-- {
			s[j], s[j-1] = s[j-1], s[j]
		}
	}
}

func newUUID() string { return randHex(16) }

func randHex(n int) string {
	b := make([]byte, n)
	if _, err := rand.Read(b); err != nil {
		return strings.Repeat("0", n*2)
	}
	return hex.EncodeToString(b)
}

func itoa(v int) string     { return fmt.Sprintf("%d", v) }
func itoa64(v int64) string { return fmt.Sprintf("%d", v) }

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	// 按字节截断可能切碎多字节字符，回退到字符边界，避免产出非法 UTF-8。
	for n > 0 && !utf8.RuneStart(s[n]) {
		n--
	}
	return s[:n] + "..."
}

func atoi(s string) int {
	n := 0
	for _, ch := range s {
		if ch < '0' || ch > '9' {
			return n
		}
		n = n*10 + int(ch-'0')
	}
	return n
}

func extractCode(data json.RawMessage) string {
	var obj struct {
		Code string `json:"code"`
	}
	if err := json.Unmarshal(data, &obj); err == nil {
		return obj.Code
	}
	return ""
}

func toInt(v any, def int) int {
	switch t := v.(type) {
	case float64:
		return int(t)
	case int:
		return t
	case string:
		return atoi(t)
	case json.Number:
		if n, err := t.Int64(); err == nil {
			return int(n)
		}
	}
	return def
}