package agent

import (
	"bufio"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

// ---------------------------------------------------------------------------
// Scenario-driven fake pi binary.
//
// The test writes a scenario file (JSONL) and sets AGENTTEAM_TEST_PI_SCENARIO in the
// wrapper script. Each line is a piFakeScript that handles one stdin command:
//
//	{"cmd":"prompt","resp":true,"events":[{frame},...]}
//	{"cmd":"steer","resp":true,"events":[...]}
//	{"cmd":"abort","resp":true,"events":[...]}
//	{"cmd":"get_state","resp":true,"data":{...}}
//
// Source-faithful failures use resp=false plus error and never include data.
//
// Scripts are consumed in order; an unscripted command gets a generic success
// response. Event frames are emitted verbatim on stdout (source-faithful).
// ---------------------------------------------------------------------------

type piFakeScript struct {
	Cmd     string           `json:"cmd"`
	Resp    *bool            `json:"resp,omitempty"`
	Error   string           `json:"error,omitempty"`
	Data    map[string]any   `json:"data,omitempty"`
	RawData json.RawMessage  `json:"rawData,omitempty"`
	Events  []map[string]any `json:"events,omitempty"`
	// WaitFor is an absolute path to a barrier file. The fake pauses (without
	// time.Sleep) until the file exists before emitting Events, enabling a
	// controllable delayed acknowledgement after the RPC response.
	WaitFor string `json:"waitFor,omitempty"`
	// WaitBeforeSettled gates ONLY the final agent_settled event behind WaitFor
	// (the preceding events flow immediately). This lets a test race a
	// generation's settlement against a replacement admission.
	WaitBeforeSettled bool `json:"waitBeforeSettled,omitempty"`
}

func writeFakePiRPCModeBinary(t *testing.T, scenario []piFakeScript) (string, string) {
	t.Helper()
	exe, err := os.Executable()
	require.NoError(t, err)

	dir := t.TempDir()
	scenarioPath := filepath.Join(dir, "scenario.jsonl")
	var sb strings.Builder
	for _, s := range scenario {
		data, _ := json.Marshal(s)
		sb.Write(data)
		sb.WriteByte('\n')
	}
	require.NoError(t, os.WriteFile(scenarioPath, []byte(sb.String()), 0o600))

	bin := filepath.Join(dir, "pi")
	if runtime.GOOS == "windows" {
		bin += ".cmd"
		script := fmt.Sprintf("@echo off\r\nset AGENTTEAM_TEST_FAKE_PI=1\r\nset AGENTTEAM_TEST_PI_SCENARIO=%s\r\n\"%s\" -test.run=TestPiRPCFakeHelper -- %%*\r\n", scenarioPath, exe)
		require.NoError(t, os.WriteFile(bin, []byte(script), 0o755))
	} else {
		script := fmt.Sprintf("#!/bin/sh\nAGENTTEAM_TEST_FAKE_PI=1 AGENTTEAM_TEST_PI_SCENARIO=%s exec %s -test.run=TestPiRPCFakeHelper -- \"$@\"\n", testShellQuote(scenarioPath), testShellQuote(exe))
		require.NoError(t, os.WriteFile(bin, []byte(script), 0o755))
	}
	return bin, scenarioPath
}

// TestPiRPCFakeHelper is the fake pi binary entrypoint.
func TestPiRPCFakeHelper(t *testing.T) {
	if os.Getenv("AGENTTEAM_TEST_FAKE_PI") != "1" {
		return
	}
	runFakePiRPCMode()
	os.Exit(0)
}

func runFakePiRPCMode() {
	lifecycleDir := strings.TrimSpace(os.Getenv("AGENTTEAM_TEST_PI_LIFECYCLE_DIR"))
	var lifecyclePath string
	if lifecycleDir != "" {
		_ = os.MkdirAll(lifecycleDir, 0o700)
		appendPiFakeLifecycleRecord(filepath.Join(lifecycleDir, "starts"), fmt.Sprintf("%d", os.Getpid()))
		matches, _ := filepath.Glob(filepath.Join(lifecycleDir, "active-*"))
		if len(matches) > 0 {
			_ = os.WriteFile(filepath.Join(lifecycleDir, "overlap"), []byte("first process was not reaped before fallback"), 0o600)
		}
		lifecyclePath = filepath.Join(lifecycleDir, fmt.Sprintf("active-%d", os.Getpid()))
		_ = os.WriteFile(lifecyclePath, []byte("active"), 0o600)
		defer os.Remove(lifecyclePath)
	}
	scripts := loadFakePiScenario(os.Getenv("AGENTTEAM_TEST_PI_SCENARIO"))
	scriptIdx := 0
	scanner := bufio.NewScanner(os.Stdin)
	scanner.Buffer(make([]byte, 64*1024), piRPCMaxLineBytes)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" {
			continue
		}
		var cmd struct {
			Type string `json:"type"`
			ID   int    `json:"id"`
		}
		if json.Unmarshal([]byte(line), &cmd) != nil {
			continue
		}
		if lifecycleDir != "" {
			appendPiFakeLifecycleRecord(filepath.Join(lifecycleDir, "commands"), cmd.Type)
		}
		script, found := nextFakePiScript(scripts, &scriptIdx, cmd.Type)
		if !found {
			success := true
			respondFake(cmd.ID, cmd.Type, &success, "", nil, nil)
			// An unscripted compact still emits a normal compaction_start/end so
			// the Process's idle-compaction completes (Pi does this in reality).
			if cmd.Type == "compact" {
				writeFakePiLine(piFakeCompactionStart("manual"))
				writeFakePiLine(piFakeCompactionEnd("manual"))
			}
			continue
		}
		if script.Resp != nil {
			respondFake(cmd.ID, cmd.Type, script.Resp, script.Error, script.Data, script.RawData)
		}
		if strings.TrimSpace(script.WaitFor) != "" && !script.WaitBeforeSettled {
			waitForFakeBarrier(script.WaitFor)
		}
		if script.WaitBeforeSettled {
			// Emit all events except the final agent_settled, then gate it behind
			// the barrier.
			pre := script.Events
			last := len(script.Events) - 1
			if last >= 0 {
				for _, evt := range script.Events[:last] {
					writeFakePiLine(evt)
				}
				pre = script.Events[last:]
			}
			if strings.TrimSpace(script.WaitFor) != "" {
				waitForFakeBarrier(script.WaitFor)
			}
			for _, evt := range pre {
				writeFakePiLine(evt)
			}
		} else {
			for _, evt := range script.Events {
				writeFakePiLine(evt)
			}
		}
	}
}

func appendPiFakeLifecycleRecord(path, value string) {
	file, err := os.OpenFile(path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o600)
	if err != nil {
		return
	}
	_, _ = fmt.Fprintln(file, value)
	_ = file.Close()
}

func loadFakePiScenario(path string) []piFakeScript {
	if path == "" {
		return nil
	}
	data, err := os.ReadFile(path)
	if err != nil {
		return nil
	}
	var scripts []piFakeScript
	for _, line := range strings.Split(strings.TrimSpace(string(data)), "\n") {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		var s piFakeScript
		if json.Unmarshal([]byte(line), &s) == nil {
			scripts = append(scripts, s)
		}
	}
	return scripts
}

func nextFakePiScript(scripts []piFakeScript, idx *int, cmdType string) (piFakeScript, bool) {
	for *idx < len(scripts) {
		candidate := scripts[*idx]
		if candidate.Cmd != cmdType {
			// Do NOT advance past a script for a different command: an unscripted
			// command (e.g. compact from idle compaction) must not consume a later
			// script meant for a future matching command.
			return piFakeScript{}, false
		}
		*idx++
		return candidate, true
	}
	return piFakeScript{}, false
}

func respondFake(id int, command string, success *bool, errorText string, data map[string]any, rawData json.RawMessage) {
	payload := map[string]any{
		"type":    "response",
		"command": command,
	}
	if success != nil {
		payload["success"] = *success
	}
	if id != 0 {
		payload["id"] = id
	}
	if errorText != "" {
		payload["error"] = errorText
	}
	if data != nil && (success == nil || *success) {
		payload["data"] = data
	}
	if len(rawData) > 0 && (success == nil || *success) {
		payload["data"] = json.RawMessage(rawData)
	}
	writeFakePiLine(payload)
}

// waitForFakeBarrier blocks until the barrier file exists, polling with
// time.After (never time.Sleep so the fixed-wait guardrail stays clean).
func waitForFakeBarrier(path string) {
	deadline := time.After(30 * time.Second)
	for {
		if _, err := os.Stat(path); err == nil {
			return
		}
		select {
		case <-deadline:
			return
		case <-time.After(10 * time.Millisecond):
		}
	}
}

func writeFakePiLine(payload any) {
	data, _ := json.Marshal(payload)
	_, _ = fmt.Fprintln(os.Stdout, string(data))
}

// --- Source-faithful frame builders for tests ---

func piFakeUserMessage(id, markerText string) map[string]any {
	content := []map[string]any{{"type": "text", "text": markerText}}
	return map[string]any{
		"type":    "message_start",
		"message": map[string]any{"role": "user", "content": content, "timestamp": 1786994272000},
	}
}

func piFakeUserMessageEnd(id, markerText string) map[string]any {
	content := []map[string]any{{"type": "text", "text": markerText}}
	return map[string]any{
		"type":    "message_end",
		"message": map[string]any{"role": "user", "content": content, "timestamp": 1786994272000},
	}
}

func piFakePlainUserMessage(markerText string) map[string]any {
	return map[string]any{
		"type":    "message_start",
		"message": map[string]any{"role": "user", "content": markerText, "timestamp": 1786994272000},
	}
}

func piFakePlainUserMessageEnd(markerText string) map[string]any {
	return map[string]any{
		"type":    "message_end",
		"message": map[string]any{"role": "user", "content": markerText, "timestamp": 1786994272000},
	}
}

func piFakeAssistantMessage(content []map[string]any, usage map[string]any) map[string]any {
	msg := map[string]any{
		"role":       "assistant",
		"content":    content,
		"provider":   "openai-codex",
		"model":      "gpt-5.5",
		"api":        "openai-codex-responses",
		"stopReason": "stop",
		"timestamp":  1786994273000,
	}
	if usage != nil {
		msg["usage"] = usage
	}
	return msg
}

func piFakeAssistantStart(content []map[string]any, usage map[string]any) map[string]any {
	return map[string]any{"type": "message_start", "message": piFakeAssistantMessage(content, usage)}
}

func piFakeAssistantEnd(content []map[string]any, usage map[string]any) map[string]any {
	return map[string]any{"type": "message_end", "message": piFakeAssistantMessage(content, usage)}
}

func piFakeTextDelta(idx int, delta string, usage map[string]any) map[string]any {
	evt := map[string]any{"type": "text_delta", "contentIndex": idx, "delta": delta}
	return map[string]any{"type": "message_update", "usage": usage, "assistantMessageEvent": evt}
}

func piFakeTextEnd(idx int, content string, usage map[string]any) map[string]any {
	evt := map[string]any{"type": "text_end", "contentIndex": idx, "content": content}
	return map[string]any{"type": "message_update", "usage": usage, "assistantMessageEvent": evt}
}

func piFakeThinkingDelta(idx int, delta string, usage map[string]any) map[string]any {
	evt := map[string]any{"type": "thinking_delta", "contentIndex": idx, "delta": delta}
	return map[string]any{"type": "message_update", "usage": usage, "assistantMessageEvent": evt}
}

func piFakeThinkingEnd(idx int, content string, usage map[string]any) map[string]any {
	evt := map[string]any{"type": "thinking_end", "contentIndex": idx, "content": content}
	return map[string]any{"type": "message_update", "usage": usage, "assistantMessageEvent": evt}
}

func piFakeToolCallEnd(idx int, toolCallID, name string, args map[string]any, usage map[string]any) map[string]any {
	evt := map[string]any{
		"type":         "toolcall_end",
		"contentIndex": idx,
		"toolCall":     map[string]any{"id": toolCallID, "name": name, "arguments": args},
	}
	return map[string]any{"type": "message_update", "usage": usage, "assistantMessageEvent": evt}
}

func piFakeToolExecutionStart(toolCallID, name string, args map[string]any) map[string]any {
	return map[string]any{"type": "tool_execution_start", "toolCallId": toolCallID, "toolName": name, "args": args}
}

func piFakeToolExecutionUpdate(toolCallID, name string, args, partialResult map[string]any) map[string]any {
	return map[string]any{"type": "tool_execution_update", "toolCallId": toolCallID, "toolName": name, "args": args, "partialResult": partialResult}
}

func piFakeToolExecutionEnd(toolCallID, name string, result map[string]any, isError bool) map[string]any {
	return map[string]any{"type": "tool_execution_end", "toolCallId": toolCallID, "toolName": name, "result": result, "isError": isError}
}

func piFakeTurnStart() map[string]any {
	return map[string]any{"type": "turn_start"}
}

func piFakeTurnEnd(content []map[string]any, usage map[string]any) map[string]any {
	evt := map[string]any{
		"type":        "turn_end",
		"message":     piFakeAssistantMessage(content, usage),
		"toolResults": []any{},
	}
	return evt
}

func piFakeAgentStart() map[string]any {
	return map[string]any{"type": "agent_start"}
}

func piFakeAgentEnd(willRetry bool, messages ...map[string]any) map[string]any {
	evt := map[string]any{"type": "agent_end", "willRetry": willRetry}
	if messages != nil {
		evt["messages"] = messages
	}
	return evt
}

func piFakeAgentSettled() map[string]any {
	return map[string]any{"type": "agent_settled"}
}

func piFakeAutoRetryStart(attempt, maxAttempts int) map[string]any {
	return map[string]any{"type": "auto_retry_start", "attempt": attempt, "maxAttempts": maxAttempts, "delayMs": 100, "errorMessage": "529 overloaded"}
}

func piFakeAutoRetryEnd(success bool, attempt int) map[string]any {
	return map[string]any{"type": "auto_retry_end", "success": success, "attempt": attempt}
}

func piFakeQueueUpdate(steering, followUp []string) map[string]any {
	return map[string]any{"type": "queue_update", "steering": steering, "followUp": followUp}
}

func piFakeCompactionStart(reason string) map[string]any {
	return map[string]any{"type": "compaction_start", "reason": reason}
}

func piFakeCompactionEnd(reason string) map[string]any {
	return map[string]any{"type": "compaction_end", "reason": reason, "result": map[string]any{"summary": "ok"}, "aborted": false, "willRetry": false}
}

// withWaitFor attaches a barrier path to a script.
func (s piFakeScript) withWaitFor(path string) piFakeScript {
	s.WaitFor = path
	return s
}

// withWaitBeforeSettled gates only the final agent_settled behind the barrier.
func (s piFakeScript) withWaitBeforeSettled(path string) piFakeScript {
	s.WaitFor = path
	s.WaitBeforeSettled = true
	return s
}

func piFakeCompactionEndAborted() map[string]any {
	return map[string]any{"type": "compaction_end", "reason": "manual", "result": nil, "aborted": true, "willRetry": false}
}

func piFakeCompactionEndRetry() map[string]any {
	return map[string]any{"type": "compaction_end", "reason": "overflow", "result": nil, "aborted": false, "willRetry": true}
}

func piFakeCompactionEndFailed() map[string]any {
	return map[string]any{"type": "compaction_end", "reason": "threshold", "result": nil, "aborted": false, "willRetry": false, "errorMessage": "provider overloaded"}
}

func piFakeCompactionEndTooSmall() map[string]any {
	return map[string]any{"type": "compaction_end", "reason": "manual", "result": nil, "aborted": false, "willRetry": false, "errorMessage": "Compaction failed: Nothing to compact (session too small)"}
}

func piFakePromptScript(events ...map[string]any) piFakeScript {
	success := true
	return piFakeScript{Cmd: "prompt", Resp: &success, Events: events}
}

func piFakeSteerScript(events ...map[string]any) piFakeScript {
	success := true
	return piFakeScript{Cmd: "steer", Resp: &success, Events: events}
}

func piFakeAbortScript(events ...map[string]any) piFakeScript {
	success := true
	return piFakeScript{Cmd: "abort", Resp: &success, Events: events}
}

func piFakeStateScript(isStreaming bool) piFakeScript {
	success := true
	return piFakeScript{Cmd: "get_state", Resp: &success, Data: map[string]any{"isStreaming": isStreaming, "isCompacting": false, "messageCount": 2, "sessionId": "test-session"}}
}

func piFakeEntriesScript() piFakeScript {
	success := true
	return piFakeScript{Cmd: "get_entries", Resp: &success, Data: map[string]any{"entries": []any{}, "leafId": "entry-leaf"}}
}

func piFakeTextBlock(text string) map[string]any {
	return map[string]any{"type": "text", "text": text}
}

func piFakeThinkingBlock(text string) map[string]any {
	return map[string]any{"type": "thinking", "thinking": text}
}

func piFakeToolCallBlock(id, name string, args map[string]any) map[string]any {
	return map[string]any{"type": "toolCall", "id": id, "name": name, "arguments": args}
}

func piFakeUsage(input, output, cacheRead, total int) map[string]any {
	return map[string]any{"input": input, "output": output, "cacheRead": cacheRead, "cacheWrite": 0, "totalTokens": total}
}