package agent

import (
	"io"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
	"github.com/tidbcloud/AgentTeam/desktop/internal/promptcompiler"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

// CodexAppServerDriver drives the Codex App Server runtime via its JSON-RPC
// stdin/stdout protocol, translating between the app-server's request/response
// model and agentteam's drivercontract.Driver interface.
type CodexAppServerDriver struct {
	mu                    sync.Mutex
	transportWriteMu      sync.Mutex
	primaryTurnDispatchMu sync.Mutex
	tokenUsageMu          sync.Mutex
	turnEventMu           sync.Mutex
	transportGenerationMu sync.RWMutex
	transportResetMu      sync.Mutex

	launchGeneration uint64

	codexRPCState
	codexLaunchState
	codexThreadState
}

type codexRPCState struct {
	requestID        int
	pending          map[int]pendingAppServerRequest
	serverRequests   map[string]inFlightServerRequest
	rawStdin         io.WriteCloser
	translatedLines  chan string
	translatedOutput *codexTranslatedOutput
	transportFailure *codexAppServerTransportFailure
}

type pendingAppServerRequest struct {
	generation uint64
	responseCh chan appServerEnvelope
}

type codexLaunchState struct {
	reasoningEffort  string
	reasoningSummary string
	showRawReasoning bool
	smallModel       string
	// launch* is the immutable per-transport configuration snapshot used by
	// later notLoaded and recovery resumes. Keep it separate from
	// effectiveModel: app-server notifications may change the latter without
	// changing AgentTeam's launch-time pinning policy.
	launchPermissionMode         string
	launchWorkDir                string
	launchModel                  string
	launchPromptCarrier          *promptcompiler.CompiledPrompt
	effectiveModel               string
	sessionResumeFallback        *drivercontract.SessionResumeFallback
	proactiveCompactionThreshold uint64
	mcpConfigs                   map[string]types.NamedMCPConfig
	trustedMCPToolReadOnly       map[string]bool
	httpClient                   *http.Client
}

const (
	codexCompactionSourceRuntimeNative      = "runtime_native"
	codexCompactionSourceAgentTeamEmergency = "agentteam_emergency"
	codexCompactionSourceLoopProactive      = "agentteam_proactive"
	codexCompactionSourceUnknown            = "unknown"
)

type codexContextCompactionItemFence struct {
	threadID        string
	turnID          string
	itemID          string
	pairingConflict bool
	before          appServerThreadTokenUsage
	during          appServerThreadTokenUsage
}

type codexContextCompactionBoundary struct {
	launchGeneration uint64
	threadID         string
	turnID           string
	itemID           string
	source           string
	causeTurnID      string
	before           appServerThreadTokenUsage
	after            appServerThreadTokenUsage
}

type codexThreadState struct {
	threadID                    string
	threadStatus                appServerThreadStatus
	activeTurnID                string
	activeTurnNonSteerableKind  string
	completedTurnCount          int
	completedTurnIDs            map[string]struct{}
	terminalErrorTurnEndIDs     map[string]struct{}
	streamedItemIDs             map[string]struct{}
	reasoningSections           map[string]reasoningSectionState
	toolActivities              map[string]toolActivityRef
	lastTokenUsage              appServerThreadTokenUsage
	lastTokenUsageByThread      map[string]appServerThreadTokenUsage
	pendingTokenUsageByThread   map[string][]pendingAppServerThreadTokenUsage
	modelByThread               map[string]string
	modelByThreadTurn           map[string]string
	modelHydrationInFlight      map[string]uint64
	modelHydrationBlocked       map[string]struct{}
	compactQueued               bool
	queuedCompactionCauseTurnID string
	compactionInProgress        bool
	compactionTurnID            string
	compactionCauseTurnID       string
	compactionSource            string
	// Codex publishes thread/status/changed idle before turn/completed. While a
	// compaction turn is active, retain that pre-terminal status until the
	// matching turn id completes; ordinary turns derive their boundary directly
	// from turn/completed and do not use this latch.
	compactionTerminalStatus         appServerThreadStatus
	contextCompactionItemFence       codexContextCompactionItemFence
	contextCompactionBoundary        codexContextCompactionBoundary
	lastContextCompactionCompletion  string
	lastContextCompactionCompletedAt time.Time
	lastActiveInterrupt              activeInterruptSnapshot
}

func (d *CodexAppServerDriver) resetAppServerStateForLaunchLocked(cfg drivercontract.LaunchConfig, stdinPipe io.WriteCloser) {
	d.launchGeneration++
	translatedOutput := newCodexTranslatedOutput(
		codexTranslatedOutputQueueEntries,
		codexTranslatedOutputMaxLineBytes,
		codexTranslatedOutputByteBudget,
	)
	d.codexRPCState = codexRPCState{
		pending:          make(map[int]pendingAppServerRequest),
		rawStdin:         stdinPipe,
		translatedLines:  translatedOutput.lines,
		translatedOutput: translatedOutput,
	}
	d.codexLaunchState = codexLaunchState{
		reasoningEffort:              strings.TrimSpace(cfg.ReasoningEffort),
		reasoningSummary:             codexReasoningSummary(cfg.Env),
		showRawReasoning:             codexShowRawReasoningFromEnv(cfg.Env),
		smallModel:                   strings.TrimSpace(cfg.SmallModel),
		launchPermissionMode:         strings.TrimSpace(cfg.PermissionMode),
		launchWorkDir:                strings.TrimSpace(cfg.WorkDir),
		launchModel:                  cfg.Model,
		launchPromptCarrier:          cfg.CompiledPrompt,
		effectiveModel:               strings.TrimSpace(cfg.Model),
		proactiveCompactionThreshold: resolveProactiveIdleCompactionThreshold(cfg),
		mcpConfigs:                   namedMCPConfigsByName(cfg.MCPs),
		trustedMCPToolReadOnly:       make(map[string]bool),
		httpClient:                   &http.Client{Timeout: 300 * time.Second},
	}
	d.codexThreadState = codexThreadState{
		completedTurnIDs:          make(map[string]struct{}),
		terminalErrorTurnEndIDs:   make(map[string]struct{}),
		streamedItemIDs:           make(map[string]struct{}),
		reasoningSections:         make(map[string]reasoningSectionState),
		lastTokenUsageByThread:    make(map[string]appServerThreadTokenUsage),
		pendingTokenUsageByThread: make(map[string][]pendingAppServerThreadTokenUsage),
		modelByThread:             make(map[string]string),
		modelByThreadTurn:         make(map[string]string),
		modelHydrationInFlight:    make(map[string]uint64),
		modelHydrationBlocked:     make(map[string]struct{}),
	}
}