package codex

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"hash/fnv"
	"io"
	"log/slog"
	"math"
	"net"
	"net/http"
	"net/url"
	"runtime"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/apr/types"
)

const providerName = "codex"

// Codex implements types.LLMProvider, types.OpenAICompatible, and types.ResponsesCompatible
// for the OpenAI Codex/ChatGPT backend.
type Codex struct {
	mode                string // "chatgpt", "api_key", or "passthrough"
	baseURL             string
	backendName         string
	backendID           int
	active              bool
	activeSet           bool
	weight              float64
	allowedAPIKeys      map[string]struct{}
	requiresOAuthBeta   bool
	customPrompt        string
	customPromptModels  []string
	tokenManager        *TokenManager
	apiKey              string
	modelCache          *types.ModelCache
	httpClient          *http.Client
	tokenHTTPClient     *http.Client
	lifecycleContext    context.Context
	authBackends        []*Codex
	fallbacks           []*Codex
	adaptiveWeights     bool
	selectorMu          sync.Mutex
	authSWRR            map[*Codex]float64
	fallbackSWRR        map[*Codex]float64
	authEffective       map[*Codex]float64
	fallbackEffective   map[*Codex]float64
	nextAuth            uint64
	nextFallback        uint64
	weightDiffThreshold float64
	lastStatus          atomic.Int64
	lastError           atomic.Value
}

var (
	_ types.LLMProvider         = (*Codex)(nil)
	_ types.OpenAICompatible    = (*Codex)(nil)
	_ types.ResponsesCompatible = (*Codex)(nil)
)

// Name returns the canonical name of the provider.
func (c *Codex) Name() string { return providerName }

// Models returns the cached list of available models.
func (c *Codex) Models() []types.Model {
	return c.modelCache.Models()
}

// ChatCompletion translates a Chat Completions request to a Responses API request,
// forwards it to the Codex backend, and translates the response back.
func (c *Codex) ChatCompletion(ctx context.Context, headers http.Header, jsonData []byte, opts ...types.CompletionOption) (*http.Response, error) {
	clientStream := extractStreamField(jsonData)

	codexJSON, nameMap, err := translateRequest(jsonData)
	if err != nil {
		return nil, fmt.Errorf("failed to translate request: %w", err)
	}
	codexJSON = c.injectCustomPrompt(codexJSON)

	resp, err := c.forward(ctx, "/responses", headers, codexJSON)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusOK {
		if clientStream {
			return streamErrorResponse(resp)
		}
		return jsonErrorResponseFromUpstream(resp)
	}

	if clientStream {
		resp.Body = newSSETranslator(resp.Body, nameMap)
		resp.Header.Set("Content-Type", "text/event-stream")
	} else {
		resp, err = translateNonStreamingResponse(resp, nameMap)
		if err != nil {
			return newJSONErrorResponse(http.StatusBadGateway, "No response.completed event received"), nil
		}
	}

	return resp, nil
}

// CreateResponse forwards a Responses API request to the Codex backend.
func (c *Codex) CreateResponse(ctx context.Context, headers http.Header, jsonData []byte) (*http.Response, error) {
	jsonData, injectedPrompt := c.injectCustomPromptWithResult(jsonData)
	resp, err := c.forward(ctx, "/responses", headers, jsonData)
	if err != nil || resp == nil || resp.Body == nil || injectedPrompt == "" || resp.StatusCode >= http.StatusBadRequest {
		return resp, err
	}

	resp.Body = newPromptSanitizingReadCloser(resp.Body, injectedPrompt, isEventStreamResponse(resp.Header))
	resp.Header.Del("Content-Length")
	resp.ContentLength = -1
	return resp, nil
}

func (c *Codex) forward(ctx context.Context, path string, headers http.Header, jsonData []byte) (*http.Response, error) {
	clientAPIKey := extractClientAPIKey(headers)
	sessionID := extractSessionID(headers)
	candidates := c.authCandidates(clientAPIKey, sessionID)
	fallbacks := c.fallbackCandidates(clientAPIKey)
	var lastErr error

	for i, candidate := range candidates {
		resp, err := candidate.forwardOnce(ctx, candidate.baseURL+path, headers, jsonData)
		if err != nil {
			c.recordAdaptiveFailure(candidate, 0)
			lastErr = err
			continue
		}
		failover, restored := shouldFailoverResponse(resp)
		if !failover {
			candidate.recordBackendStatus(restored.StatusCode)
			c.recordAdaptiveSuccess(candidate)
			return restored, nil
		}
		candidate.recordBackendStatus(restored.StatusCode)
		c.recordAdaptiveFailure(candidate, restored.StatusCode)
		slog.Warn("Codex backend triggered failover",
			"backend", candidate.backendLabel(),
			"status", restored.StatusCode,
			"path", path,
		)
		if i == len(candidates)-1 && len(fallbacks) == 0 {
			return restored, nil
		}
		restored.Body.Close()
	}

	for i, fallback := range fallbacks {
		resp, err := fallback.forward(ctx, path, headers, jsonData)
		if err != nil {
			lastErr = err
			continue
		}
		failover, restored := shouldFailoverResponse(resp)
		if !failover {
			fallback.recordBackendStatus(restored.StatusCode)
			c.recordAdaptiveFallbackSuccess(fallback)
			return restored, nil
		}
		fallback.recordBackendStatus(restored.StatusCode)
		c.recordAdaptiveFallbackFailure(fallback, restored.StatusCode)
		slog.Warn("Codex fallback backend triggered failover",
			"backend", fallback.backendLabel(),
			"status", restored.StatusCode,
			"path", path,
		)
		if i == len(fallbacks)-1 {
			return restored, nil
		}
		restored.Body.Close()
	}
	if lastErr != nil {
		return nil, lastErr
	}
	if len(candidates) == 0 && len(fallbacks) == 0 {
		return newJSONErrorResponse(http.StatusForbidden, "No Codex backend is active for this API key"), nil
	}
	return newJSONErrorResponse(http.StatusServiceUnavailable, "All Codex backends failed over"), nil
}

func (c *Codex) authCandidates(clientAPIKey, sessionID string) []*Codex {
	all := make([]*Codex, 0, 1+len(c.authBackends))
	all = append(all, c)
	all = append(all, c.authBackends...)
	return c.orderBackends(all, clientAPIKey, sessionID, &c.nextAuth, c.authSWRR, c.authEffective)
}

func (c *Codex) fallbackCandidates(clientAPIKey string) []*Codex {
	if len(c.fallbacks) == 0 {
		return nil
	}
	return c.orderBackends(c.fallbacks, clientAPIKey, "", &c.nextFallback, c.fallbackSWRR, c.fallbackEffective)
}

func (c *Codex) orderBackends(backends []*Codex, clientAPIKey, sessionID string, counter *uint64, swrr map[*Codex]float64, effective map[*Codex]float64) []*Codex {
	eligible := make([]*Codex, 0, len(backends))
	for _, backend := range backends {
		if backend != nil && backend.isEligible(clientAPIKey) {
			eligible = append(eligible, backend)
		}
	}
	if len(eligible) <= 1 {
		return eligible
	}

	selected := c.selectInitialBackend(eligible, sessionID, counter, swrr, effective)
	if selected == nil {
		return eligible
	}

	ordered := make([]*Codex, 0, len(eligible))
	ordered = append(ordered, selected)
	rest := make([]*Codex, 0, len(eligible)-1)
	for _, backend := range eligible {
		if backend != selected {
			rest = append(rest, backend)
		}
	}
	sort.SliceStable(rest, func(i, j int) bool {
		ei := c.effectiveWeight(rest[i], effective)
		ej := c.effectiveWeight(rest[j], effective)
		if ei != ej {
			return ei > ej
		}
		return rest[i].backendLabel() < rest[j].backendLabel()
	})
	ordered = append(ordered, rest...)
	return ordered
}

func (c *Codex) selectInitialBackend(backends []*Codex, sessionID string, counter *uint64, swrr map[*Codex]float64, effective map[*Codex]float64) *Codex {
	var totalWeight float64
	for _, backend := range backends {
		w := c.effectiveWeight(backend, effective)
		if w > 0 {
			totalWeight += w
		}
	}
	if totalWeight <= 0 {
		idx := int((atomic.AddUint64(counter, 1) - 1) % uint64(len(backends)))
		return backends[idx]
	}

	// Session affinity: hash-based deterministic selection.
	if sessionID != "" {
		return c.selectBySessionHash(backends, sessionID, effective, totalWeight)
	}

	// Dominant-weight short-circuit: if top weight exceeds second by threshold.
	if len(backends) >= 2 && c.weightDiffThreshold > 0 {
		sorted := make([]*Codex, len(backends))
		copy(sorted, backends)
		sort.SliceStable(sorted, func(i, j int) bool {
			return c.effectiveWeight(sorted[i], effective) > c.effectiveWeight(sorted[j], effective)
		})
		top := c.effectiveWeight(sorted[0], effective)
		second := c.effectiveWeight(sorted[1], effective)
		if top-second > c.weightDiffThreshold {
			return sorted[0]
		}
	}

	// Standard SWRR.
	if swrr == nil {
		idx := int((atomic.AddUint64(counter, 1) - 1) % uint64(len(backends)))
		return backends[idx]
	}
	c.selectorMu.Lock()
	defer c.selectorMu.Unlock()

	for existing := range swrr {
		found := false
		for _, backend := range backends {
			if existing == backend && c.effectiveWeight(backend, effective) > 0 {
				found = true
				break
			}
		}
		if !found {
			delete(swrr, existing)
		}
	}

	var selected *Codex
	for _, backend := range backends {
		w := c.effectiveWeight(backend, effective)
		if w <= 0 {
			continue
		}
		swrr[backend] += w
		if selected == nil || swrr[backend] > swrr[selected] {
			selected = backend
		}
	}
	if selected != nil {
		swrr[selected] -= totalWeight
	}
	return selected
}

func (c *Codex) selectBySessionHash(backends []*Codex, sessionID string, effective map[*Codex]float64, totalWeight float64) *Codex {
	h := fnv.New32a()
	h.Write([]byte(sessionID))
	slot := float64(h.Sum32()) / float64(math.MaxUint32) * totalWeight

	cumulative := 0.0
	for _, backend := range backends {
		cumulative += c.effectiveWeight(backend, effective)
		if slot < cumulative {
			return backend
		}
	}
	return backends[0]
}

func (c *Codex) effectiveWeight(backend *Codex, effective map[*Codex]float64) float64 {
	if ew, ok := effective[backend]; ok && ew > 0 {
		return ew
	}
	return backend.weight
}

// recordAdaptiveSuccess increases effective weight on a successful request.
func (c *Codex) recordAdaptiveSuccess(backend *Codex) {
	if !c.adaptiveWeights || backend == nil || backend.weight <= 0 {
		return
	}
	c.selectorMu.Lock()
	defer c.selectorMu.Unlock()

	ew, ok := c.authEffective[backend]
	if !ok {
		ew = backend.weight
	}
	ew = math.Min(ew*1.1, backend.weight)
	if ew < 1 {
		ew = 1
	}
	c.authEffective[backend] = ew
}

// recordAdaptiveFailure decreases effective weight based on the error status code.
func (c *Codex) recordAdaptiveFailure(backend *Codex, statusCode int) {
	if !c.adaptiveWeights || backend == nil || backend.weight <= 0 {
		return
	}
	c.selectorMu.Lock()
	defer c.selectorMu.Unlock()

	c.adjustEffectiveWeightLocked(backend, statusCode, c.authEffective)
}

func (c *Codex) recordAdaptiveFallbackSuccess(backend *Codex) {
	if !c.adaptiveWeights || backend == nil || backend.weight <= 0 {
		return
	}
	c.selectorMu.Lock()
	defer c.selectorMu.Unlock()

	ew, ok := c.fallbackEffective[backend]
	if !ok {
		ew = backend.weight
	}
	ew = math.Min(ew*1.1, backend.weight)
	if ew < 1 {
		ew = 1
	}
	c.fallbackEffective[backend] = ew
}

func (c *Codex) recordAdaptiveFallbackFailure(backend *Codex, statusCode int) {
	if !c.adaptiveWeights || backend == nil || backend.weight <= 0 {
		return
	}
	c.selectorMu.Lock()
	defer c.selectorMu.Unlock()

	c.adjustEffectiveWeightLocked(backend, statusCode, c.fallbackEffective)
}

func (c *Codex) adjustEffectiveWeightLocked(backend *Codex, statusCode int, effective map[*Codex]float64) {
	ew, ok := c.authEffective[backend]
	if effective != nil {
		ew, ok = effective[backend]
	}
	if !ok {
		ew = backend.weight
	}
	factor := 0.8
	switch {
	case statusCode == 429:
		factor = 0.5
	case statusCode == 503:
		factor = 0.6
	}
	ew *= factor
	minWeight := math.Max(1, backend.weight*0.1)
	if ew < minWeight {
		ew = minWeight
	}
	if effective != nil {
		effective[backend] = ew
	}
}

// extractSessionID extracts session affinity key from request headers.
func extractSessionID(headers http.Header) string {
	for _, name := range []string{"Session-Id", "session_id", "Conversation-Id", "conversation_id"} {
		if v := strings.TrimSpace(headers.Get(name)); v != "" {
			return v
		}
	}
	return ""
}

func (c *Codex) isEligible(clientAPIKey string) bool {
	if c.activeSet && !c.active {
		return false
	}
	if len(c.allowedAPIKeys) == 0 {
		return true
	}
	if clientAPIKey == "" {
		return false
	}
	_, ok := c.allowedAPIKeys[clientAPIKey]
	return ok
}

func (c *Codex) backendLabel() string {
	if c.backendName != "" {
		return c.backendName
	}
	if c.backendID != 0 {
		return fmt.Sprintf("%d", c.backendID)
	}
	if c.baseURL != "" {
		return c.baseURL
	}
	return providerName
}

func (c *Codex) recordBackendStatus(status int) {
	c.lastStatus.Store(int64(status))
	c.lastError.Store("")
}

func (c *Codex) recordBackendError(err error) {
	if err == nil {
		return
	}
	c.lastError.Store(err.Error())
}

func extractClientAPIKey(headers http.Header) string {
	if headers == nil {
		return ""
	}
	if auth := strings.TrimSpace(headers.Get("Authorization")); auth != "" {
		parts := strings.SplitN(auth, " ", 2)
		if len(parts) == 2 && strings.EqualFold(parts[0], "Bearer") {
			return strings.TrimSpace(parts[1])
		}
	}
	for _, name := range []string{"x-openai-api-key", "openai-api-key", "x-api-key", "x-openai-key"} {
		if value := strings.TrimSpace(headers.Get(name)); value != "" {
			return value
		}
	}
	return ""
}

func (c *Codex) injectCustomPrompt(jsonData []byte) []byte {
	out, _ := c.injectCustomPromptWithResult(jsonData)
	return out
}

func (c *Codex) injectCustomPromptWithResult(jsonData []byte) ([]byte, string) {
	prompt := strings.TrimSpace(c.customPrompt)
	if prompt == "" {
		return jsonData, ""
	}

	var body map[string]interface{}
	if err := json.Unmarshal(jsonData, &body); err != nil {
		return jsonData, ""
	}
	if !c.customPromptApplies(body) {
		return jsonData, ""
	}

	if _, ok := body["input"]; ok {
		existing, _ := body["instructions"].(string)
		if strings.TrimSpace(existing) != "" {
			body["instructions"] = existing + "\n\n" + prompt
		} else {
			body["instructions"] = prompt
		}
		updated, err := json.Marshal(body)
		if err != nil {
			return jsonData, ""
		}
		return updated, prompt
	}

	if messages, ok := body["messages"].([]interface{}); ok {
		if len(messages) > 0 {
			if first, ok := messages[0].(map[string]interface{}); ok && first["role"] == "system" {
				switch content := first["content"].(type) {
				case string:
					if strings.TrimSpace(content) != "" {
						first["content"] = content + "\n\n" + prompt
					} else {
						first["content"] = prompt
					}
				case []interface{}:
					first["content"] = append(content, map[string]interface{}{"type": "text", "text": "\n\n" + prompt})
				default:
					first["content"] = prompt
				}
				updated, err := json.Marshal(body)
				if err != nil {
					return jsonData, ""
				}
				return updated, prompt
			}
		}
		body["messages"] = append([]interface{}{map[string]interface{}{"role": "system", "content": prompt}}, messages...)
		updated, err := json.Marshal(body)
		if err != nil {
			return jsonData, ""
		}
		return updated, prompt
	}

	return jsonData, ""
}

func (c *Codex) customPromptApplies(body map[string]interface{}) bool {
	if len(c.customPromptModels) == 0 {
		return true
	}
	model, _ := body["model"].(string)
	model = strings.ToLower(model)
	for _, filter := range c.customPromptModels {
		filter = strings.ToLower(strings.TrimSpace(filter))
		if filter != "" && strings.Contains(model, filter) {
			return true
		}
	}
	return false
}

func isEventStreamResponse(headers http.Header) bool {
	return strings.Contains(strings.ToLower(headers.Get("Content-Type")), "text/event-stream")
}

func newPromptSanitizingReadCloser(body io.ReadCloser, prompt string, stream bool) io.ReadCloser {
	if stream {
		scanner := bufio.NewScanner(body)
		scanner.Buffer(make([]byte, 0, 64*1024), 1*1024*1024)
		return &promptSanitizingSSEReadCloser{
			body:    body,
			scanner: scanner,
			prompt:  prompt,
		}
	}
	return &promptSanitizingJSONReadCloser{
		body:   body,
		prompt: prompt,
	}
}

type promptSanitizingJSONReadCloser struct {
	body   io.ReadCloser
	prompt string
	once   sync.Once
	reader *bytes.Reader
	err    error
}

func (r *promptSanitizingJSONReadCloser) Read(p []byte) (int, error) {
	r.once.Do(func() {
		body, err := io.ReadAll(r.body)
		if err != nil {
			r.err = err
			return
		}
		body = sanitizePromptJSONBytes(body, r.prompt)
		r.reader = bytes.NewReader(body)
	})
	if r.err != nil {
		return 0, r.err
	}
	return r.reader.Read(p)
}

func (r *promptSanitizingJSONReadCloser) Close() error {
	return r.body.Close()
}

type promptSanitizingSSEReadCloser struct {
	body    io.ReadCloser
	scanner *bufio.Scanner
	prompt  string
	buf     bytes.Buffer
	done    bool
}

func (r *promptSanitizingSSEReadCloser) Read(p []byte) (int, error) {
	for r.buf.Len() == 0 {
		if r.done {
			return 0, io.EOF
		}
		if !r.scanner.Scan() {
			r.done = true
			if err := r.scanner.Err(); err != nil {
				return 0, err
			}
			return 0, io.EOF
		}
		r.buf.WriteString(sanitizePromptSSELine(r.scanner.Text(), r.prompt))
		r.buf.WriteByte('\n')
	}
	return r.buf.Read(p)
}

func (r *promptSanitizingSSEReadCloser) Close() error {
	return r.body.Close()
}

func sanitizePromptSSELine(line, prompt string) string {
	trimmed := strings.TrimSpace(line)
	if !strings.HasPrefix(trimmed, "data:") || !strings.Contains(trimmed, prompt) {
		return line
	}
	data := strings.TrimSpace(trimmed[5:])
	if data == "" {
		return line
	}
	sanitized := sanitizePromptJSONBytes([]byte(data), prompt)
	if bytes.Equal(sanitized, []byte(data)) {
		return line
	}
	return "data: " + string(sanitized)
}

func sanitizePromptJSONBytes(body []byte, prompt string) []byte {
	if prompt == "" || !bytes.Contains(body, []byte(prompt)) {
		return body
	}
	var obj map[string]interface{}
	if err := json.Unmarshal(body, &obj); err != nil {
		return body
	}
	if !stripInjectedPromptFromObj(obj, prompt) {
		return body
	}
	updated, err := json.Marshal(obj)
	if err != nil {
		return body
	}
	return updated
}

var promptResponseFields = []string{"instructions", "system", "system_instruction", "systemInstruction"}

func stripInjectedPromptFromObj(obj map[string]interface{}, prompt string) bool {
	modified := false
	if stripInjectedPromptFromMap(obj, prompt) {
		modified = true
	}
	if response, ok := obj["response"].(map[string]interface{}); ok {
		if stripInjectedPromptFromMap(response, prompt) {
			modified = true
		}
	}
	return modified
}

func stripInjectedPromptFromMap(obj map[string]interface{}, prompt string) bool {
	modified := false
	for _, field := range promptResponseFields {
		value, ok := obj[field]
		if !ok {
			continue
		}
		cleaned, changed := stripInjectedPromptFromValue(value, prompt)
		if changed {
			obj[field] = cleaned
			modified = true
		}
	}
	return modified
}

func stripInjectedPromptFromValue(value interface{}, prompt string) (interface{}, bool) {
	switch typed := value.(type) {
	case string:
		return stripInjectedPromptFromString(typed, prompt)
	case []interface{}:
		cleaned := make([]interface{}, 0, len(typed))
		changed := false
		for _, item := range typed {
			block, ok := item.(map[string]interface{})
			if !ok || block["type"] != "text" {
				cleaned = append(cleaned, item)
				continue
			}
			text, ok := block["text"].(string)
			if !ok {
				cleaned = append(cleaned, item)
				continue
			}
			if strings.TrimSpace(text) == strings.TrimSpace(prompt) {
				changed = true
				continue
			}
			cleanText, textChanged := stripInjectedPromptFromString(text, prompt)
			if textChanged {
				next := make(map[string]interface{}, len(block))
				for k, v := range block {
					next[k] = v
				}
				next["text"] = cleanText
				cleaned = append(cleaned, next)
				changed = true
				continue
			}
			cleaned = append(cleaned, item)
		}
		if changed {
			return cleaned, true
		}
	}
	return value, false
}

func stripInjectedPromptFromString(value, prompt string) (string, bool) {
	original := value
	value = strings.ReplaceAll(value, "\n\n"+prompt, "")
	value = strings.ReplaceAll(value, prompt+"\n\n", "")
	if strings.TrimSpace(value) == strings.TrimSpace(prompt) {
		return "", true
	}
	value = strings.ReplaceAll(value, prompt, "")
	return value, value != original
}

func (c *Codex) forwardOnce(ctx context.Context, endpoint string, headers http.Header, jsonData []byte) (*http.Response, error) {
	req, err := http.NewRequestWithContext(ctx, "POST", endpoint, bytes.NewReader(jsonData))
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	parsed, _ := url.Parse(endpoint)
	host := ""
	if parsed != nil {
		host = parsed.Host
	}

	if shouldMinimizeHeaders(host) {
		c.setChatGPTHeaders(req, host, headers)
	} else {
		c.setStandardHeaders(req, host, headers)
	}

	resp, err := c.httpClient.Do(req)
	if err != nil {
		c.recordBackendError(err)
		return nil, err
	}

	// On upstream 401 in chatgpt mode, force-refresh the token and retry once
	if resp.StatusCode == http.StatusUnauthorized && c.tokenManager != nil {
		resp.Body.Close()

		ctx := req.Context()
		if err := c.tokenManager.ForceRefresh(ctx); err != nil {
			slog.Warn("Force token refresh on 401 failed", "error", err)
			return newJSONErrorResponse(http.StatusUnauthorized, "Upstream 401 and token refresh failed: "+err.Error()), nil
		}

		// Rebuild request with refreshed token
		req, err = http.NewRequestWithContext(ctx, "POST", endpoint, bytes.NewReader(jsonData))
		if err != nil {
			return nil, fmt.Errorf("failed to create retry request: %w", err)
		}
		if shouldMinimizeHeaders(host) {
			c.setChatGPTHeaders(req, host, headers)
		} else {
			c.setStandardHeaders(req, host, headers)
		}
		resp, err = c.httpClient.Do(req)
		if err != nil {
			c.recordBackendError(err)
			return nil, err
		}
		return resp, nil
	}

	return resp, nil
}

func shouldFailoverResponse(resp *http.Response) (bool, *http.Response) {
	if resp == nil {
		return true, resp
	}
	status := resp.StatusCode
	if status == http.StatusUnauthorized ||
		status == http.StatusForbidden ||
		status == http.StatusTooManyRequests ||
		status == http.StatusServiceUnavailable ||
		status == 529 ||
		status >= http.StatusInternalServerError {
		return true, resp
	}

	if status != http.StatusBadRequest || resp.Body == nil {
		return false, resp
	}

	body, err := io.ReadAll(resp.Body)
	resp.Body.Close()
	resp.Body = io.NopCloser(bytes.NewReader(body))
	resp.ContentLength = int64(len(body))
	if err != nil {
		return false, resp
	}

	text := strings.ToLower(string(body))
	if strings.Contains(text, "model is not supported when using codex with a chatgpt account") ||
		strings.Contains(text, "not supported when using codex with a chatgpt account") ||
		(strings.Contains(text, "gpt-5.3-codex") && strings.Contains(text, "chatgpt account") && strings.Contains(text, "not supported")) {
		return true, resp
	}
	return false, resp
}

func streamErrorResponse(resp *http.Response) (*http.Response, error) {
	status := resp.StatusCode
	body, _ := io.ReadAll(resp.Body)
	resp.Body.Close()

	msg := fmt.Sprintf("Upstream error %d: %s", status, string(body))
	payload, _ := json.Marshal(map[string]interface{}{
		"error": map[string]interface{}{
			"message": msg,
			"type":    "server_error",
		},
	})
	data := append([]byte("data: "), payload...)
	data = append(data, []byte("\n\ndata: [DONE]\n\n")...)

	return &http.Response{
		StatusCode:    http.StatusOK,
		Header:        http.Header{"Content-Type": []string{"text/event-stream"}},
		Body:          io.NopCloser(bytes.NewReader(data)),
		ContentLength: int64(len(data)),
	}, nil
}

func jsonErrorResponseFromUpstream(resp *http.Response) (*http.Response, error) {
	status := resp.StatusCode
	io.Copy(io.Discard, resp.Body)
	resp.Body.Close()
	return newJSONErrorResponse(status, fmt.Sprintf("Upstream error %d", status)), nil
}

func newJSONErrorResponse(status int, message string) *http.Response {
	payload, _ := json.Marshal(map[string]interface{}{
		"error": map[string]interface{}{
			"message": message,
			"type":    "server_error",
		},
	})
	return &http.Response{
		StatusCode:    status,
		Header:        http.Header{"Content-Type": []string{"application/json"}},
		Body:          io.NopCloser(bytes.NewReader(payload)),
		ContentLength: int64(len(payload)),
	}
}

// shouldMinimizeHeaders returns true for chatgpt.com and chat.openai.com hosts,
// which require a minimal header set to avoid rejection.
func shouldMinimizeHeaders(host string) bool {
	if host == "" {
		return false
	}
	hostname := strings.SplitN(host, ":", 2)[0]
	return strings.HasSuffix(hostname, "chatgpt.com") || strings.HasSuffix(hostname, "chat.openai.com")
}

// setChatGPTHeaders builds minimal headers for chatgpt.com requests.
// Only whitelisted headers are forwarded from the client; the rest are constructed.
func (c *Codex) setChatGPTHeaders(req *http.Request, host string, original http.Header) {
	// Whitelist of client headers to forward (stored with underscores for normalization)
	allowed := map[string]bool{
		"connection":      true,
		"content_type":    true,
		"conversation_id": true,
		"openai_beta":     true,
		"originator":      true,
		"session_id":      true,
		"user_agent":      true,
	}

	// Forward whitelisted headers from client (support dash and underscore variants)
	for name, values := range original {
		normalized := strings.ReplaceAll(strings.ToLower(name), "-", "_")
		if allowed[normalized] {
			for _, v := range values {
				req.Header.Add(name, v)
			}
		}
	}

	// Defaults (only set if client didn't provide them)
	req.Host = host
	if req.Header.Get("Connection") == "" {
		req.Header.Set("Connection", "keep-alive")
	}
	if req.Header.Get("Content-Type") == "" {
		req.Header.Set("Content-Type", "application/json")
	}
	// NOTE: Codex CLI does NOT send openai-beta header; only pass through if client sends it
	if req.Header.Get("Originator") == "" {
		req.Header.Set("Originator", "codex_cli_rs")
	}

	// Inject auth (mutually exclusive: tokenManager for chatgpt mode, apiKey for api_key mode)
	if c.tokenManager != nil {
		accessToken, accountID := c.tokenManager.GetAccessToken()
		if accessToken != "" {
			req.Header.Set("Authorization", "Bearer "+accessToken)
		}
		if accountID != "" {
			req.Header.Set("Chatgpt-Account-Id", accountID)
		}
	} else if c.apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+c.apiKey)
	}
}

// nonChatGPTExcludedHeaders are headers stripped when forwarding to non-chatgpt.com hosts.
var nonChatGPTExcludedHeaders = map[string]bool{
	"x-api-key":                   true,
	"authorization":               true,
	"host":                        true,
	"content-length":              true,
	"transfer-encoding":           true,
	"x-real-ip":                   true,
	"x-forwarded-for":             true,
	"x-forwarded-proto":           true,
	"cf-connecting-ip":            true,
	"cf-ray":                      true,
	"cf-visitor":                  true,
	"cf-ipcountry":                true,
	"cdn-agentteam":               true,
	"true-client-ip":              true,
	"x-stainless-os":              true,
	"x-stainless-arch":            true,
	"x-stainless-runtime":         true,
	"x-stainless-runtime-version": true,
}

// setStandardHeaders builds headers for non-chatgpt.com hosts.
// Client headers are forwarded except for the exclusion list; auth is injected.
func (c *Codex) setStandardHeaders(req *http.Request, host string, original http.Header) {
	// Forward all client headers except excluded ones
	for name, values := range original {
		if nonChatGPTExcludedHeaders[strings.ToLower(name)] {
			continue
		}
		for _, v := range values {
			req.Header.Add(name, v)
		}
	}

	req.Host = host
	if req.Header.Get("Connection") == "" {
		req.Header.Set("Connection", "keep-alive")
	}

	// Inject auth (mutually exclusive: tokenManager for chatgpt mode, apiKey for api_key mode)
	if c.tokenManager != nil {
		accessToken, _ := c.tokenManager.GetAccessToken()
		if accessToken != "" {
			req.Header.Set("Authorization", "Bearer "+accessToken)
		}
	} else if c.apiKey != "" {
		req.Header.Set("X-Api-Key", c.apiKey)
	}
	if c.requiresOAuthBeta {
		anthropicBeta := req.Header.Get("Anthropic-Beta")
		if !strings.Contains(anthropicBeta, "oauth-2025-04-20") {
			if anthropicBeta != "" {
				req.Header.Set("Anthropic-Beta", anthropicBeta+",oauth-2025-04-20")
			} else {
				req.Header.Set("Anthropic-Beta", "oauth-2025-04-20")
			}
		}
		req.Header.Del("X-Api-Key")
	}

	req.Header.Set("Content-Type", "application/json")

	// Inject proxy info headers (replacing stripped client headers)
	injectProxyInfoHeaders(req)
}

// proxySystemInfo is computed once at init and cached.
var proxySystemInfo struct {
	os             string
	arch           string
	runtimeName    string
	runtimeVersion string
	realIP         string
}

func init() {
	os := runtime.GOOS
	switch os {
	case "linux":
		os = "Linux"
	case "darwin":
		os = "MacOS"
	case "windows":
		os = "Windows"
	}
	arch := runtime.GOARCH
	switch arch {
	case "amd64":
		arch = "x64"
	case "arm64":
		arch = "arm64"
	}
	proxySystemInfo.os = os
	proxySystemInfo.arch = arch
	proxySystemInfo.runtimeName = "go"
	proxySystemInfo.runtimeVersion = runtime.Version()
	proxySystemInfo.realIP = detectProxyRealIP()
}

// injectProxyInfoHeaders injects proxy server identity headers.
// These replace the client headers stripped by the exclusion list.
func injectProxyInfoHeaders(req *http.Request) {
	req.Header.Set("X-Real-Ip", proxySystemInfo.realIP)
	req.Header.Set("X-Forwarded-For", proxySystemInfo.realIP)
	req.Header.Set("X-Stainless-Os", proxySystemInfo.os)
	req.Header.Set("X-Stainless-Arch", proxySystemInfo.arch)
	req.Header.Set("X-Stainless-Runtime", proxySystemInfo.runtimeName)
	req.Header.Set("X-Stainless-Runtime-Version", proxySystemInfo.runtimeVersion)
}

func detectProxyRealIP() string {
	conn, err := net.DialTimeout("udp", "8.8.8.8:80", time.Second)
	if err != nil {
		return "127.0.0.1"
	}
	defer conn.Close()
	if addr, ok := conn.LocalAddr().(*net.UDPAddr); ok && addr.IP != nil {
		return addr.IP.String()
	}
	return "127.0.0.1"
}

// TestResult holds the result of a Codex endpoint connectivity test.
type TestResult struct {
	Success    bool   `json:"success"`
	StatusCode int    `json:"status_code,omitempty"`
	DurationMS int64  `json:"duration_ms"`
	Error      string `json:"error,omitempty"`
	TargetURL  string `json:"target_url,omitempty"`
}

// TestEndpoint verifies connectivity to the Codex backend with a minimal request.
func (c *Codex) TestEndpoint(ctx context.Context, model string, extraParams map[string]interface{}) *TestResult {
	if model == "" {
		model = "codex-mini"
	}
	instructions := "Reply with a short OK for connectivity verification."
	if ep, ok := extraParams["instructions"].(string); ok && strings.TrimSpace(ep) != "" {
		instructions = strings.TrimSpace(ep)
	}

	body := map[string]interface{}{
		"model":             model,
		"instructions":      instructions,
		"input":             []interface{}{map[string]interface{}{"type": "message", "role": "user", "content": []interface{}{map[string]interface{}{"type": "input_text", "text": "."}}}},
		"stream":            false,
		"max_output_tokens": 1,
	}
	if re, ok := extraParams["reasoning_effort"].(string); ok && re != "" {
		body["reasoning"] = map[string]interface{}{"effort": re}
	}

	jsonData, err := json.Marshal(body)
	if err != nil {
		return &TestResult{Success: false, Error: err.Error(), TargetURL: c.baseURL + "/responses"}
	}

	start := time.Now()
	resp, err := c.forwardOnce(ctx, c.baseURL+"/responses", http.Header{}, jsonData)
	duration := time.Since(start)
	result := &TestResult{
		DurationMS: duration.Milliseconds(),
		TargetURL:  c.baseURL + "/responses",
	}

	if err != nil {
		result.Error = err.Error()
		return result
	}
	defer resp.Body.Close()

	result.StatusCode = resp.StatusCode
	result.Success = resp.StatusCode == http.StatusOK
	if !result.Success {
		body, _ := io.ReadAll(resp.Body)
		result.Error = fmt.Sprintf("HTTP %d: %s", resp.StatusCode, string(body))
	}
	return result
}

// correctBaseURL appends /backend-api/codex to chatgpt.com URLs if missing.
func correctBaseURL(baseURL string) string {
	parsed, err := url.Parse(baseURL)
	if err != nil {
		return baseURL
	}

	if strings.HasSuffix(parsed.Host, "chatgpt.com") && !strings.Contains(baseURL, "/backend-api/codex") {
		return strings.TrimRight(baseURL, "/") + "/backend-api/codex"
	}

	return baseURL
}
