package agent

import (
	"bufio"
	"context"
	"encoding/json"
	"io"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/runtimeconfig"
	"github.com/tidbcloud/AgentTeam/desktop/internal/consts"
)

func TestPiRPCParseLine(t *testing.T) {
	d := &PiRPCDriver{}

	tests := []struct {
		name string
		line string
		want string
	}{
		{"session_init", `{"kind":"session_init","sessionId":"s1"}`, consts.ParsedEventKindSessionInit},
		{"text", `{"kind":"text","text":"hello","sessionId":"s1"}`, consts.ParsedEventKindText},
		{"thinking", `{"kind":"thinking","text":"reasoning","sessionId":"s1"}`, consts.ParsedEventKindThinking},
		{"tool_call", `{"kind":"tool_call","toolName":"bash","sessionId":"s1"}`, consts.ParsedEventKindToolCall},
		{"tool_output", `{"kind":"tool_output","text":"out","toolName":"bash","sessionId":"s1"}`, consts.ParsedEventKindToolOutput},
		{"turn_end", `{"kind":"turn_end","sessionId":"s1","stopReason":"end_turn"}`, consts.ParsedEventKindTurnEnd},
		{"runtime_idle", `{"kind":"runtime_idle","sessionId":"s1"}`, consts.ParsedEventKindRuntimeIdle},
		{"error", `{"kind":"error","message":"boom","sessionId":"s1"}`, consts.ParsedEventKindError},
		{"usage", `{"kind":"usage","inputTokens":10,"outputTokens":5,"sessionId":"s1"}`, consts.ParsedEventKindUsage},
		{"runtime_input_surfaced", `{"kind":"runtime_input_surfaced","runtimeInputId":"msg_1","sessionId":"s1"}`, consts.ParsedEventKindRuntimeInputAcked},
		{"compacting", `{"kind":"runtime_compacting","sessionId":"s1"}`, consts.ParsedEventKindRuntimeCompacting},
		{"non_json_ignored", `not json`, ""},
		{"unknown_kind_ignored", `{"kind":"wat","sessionId":"s1"}`, ""},
		{"empty_text_ignored", `{"kind":"text","text":"","sessionId":"s1"}`, ""},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			events := d.ParseLine(tt.line)
			if tt.want == "" {
				require.Empty(t, events)
				return
			}
			require.Len(t, events, 1)
			require.Equal(t, tt.want, events[0].Kind)
		})
	}
}

func TestPiRPCEncodeStdinMessage(t *testing.T) {
	d := &PiRPCDriver{}
	data := d.EncodeStdinMessage("hello", "session")
	require.True(t, strings.HasSuffix(string(data), "\n"))
	var command map[string]any
	require.NoError(t, json.Unmarshal(data, &command))
	require.Equal(t, "prompt", command["type"])
	require.Equal(t, "hello", command["message"])
	require.Equal(t, "followUp", command["streamingBehavior"])
}

// ---------------------------------------------------------------------------
// Collection helpers
// ---------------------------------------------------------------------------

type piCollectedEvent struct {
	kind                string
	text                string
	message             string
	toolName            string
	toolCallID          string
	toolInput           string
	runtimeInputID      string
	stopReason          string
	model               string
	invocationID        string
	inputTokens         uint64
	outputTokens        uint64
	cachedInputTokens   uint64
	activeContextTokens uint64
}

// collectPiEvents reads translated events from the driver stdout until `until`
// kind is observed (or the deadline fires). It returns the ordered event list.
func collectPiEvents(t *testing.T, result drivercontract.SpawnResult, until string) []piCollectedEvent {
	return collectPiEventsN(t, result, until, 1)
}

// collectPiEventsN reads until `until` is observed `count` times.
func collectPiEventsN(t *testing.T, result drivercontract.SpawnResult, until string, count int) []piCollectedEvent {
	t.Helper()
	events := []piCollectedEvent{}
	seen := 0
	deadline := time.After(30 * time.Second)
	for ev := range piStreamEvents(t, result) {
		events = append(events, ev)
		if ev.kind == until {
			seen++
			if seen >= count {
				break
			}
		}
		select {
		case <-deadline:
			t.Fatalf("timed out collecting events (got %d, seen %d)", len(events), seen)
		default:
		}
	}
	return events
}

// piStreamEvents streams translated events from the driver stdout via a
// goroutine so readers can select on a deadline without blocking on Scan.
// The goroutine exits when the pipe closes (result.Cleanup).
func piStreamEvents(t *testing.T, result drivercontract.SpawnResult) <-chan piCollectedEvent {
	t.Helper()
	ch := make(chan piCollectedEvent, 256)
	go func() {
		defer close(ch)
		scanner := bufio.NewScanner(result.Stdout)
		scanner.Buffer(make([]byte, 64*1024), piRPCMaxLineBytes)
		for scanner.Scan() {
			line := strings.TrimSpace(scanner.Text())
			if line == "" {
				continue
			}
			var ev translatedEvent
			if json.Unmarshal([]byte(line), &ev) != nil {
				continue
			}
			collected := piCollectedEvent{
				kind:                ev.Kind,
				text:                ev.Text,
				message:             ev.Message,
				toolName:            ev.ToolName,
				toolCallID:          ev.ToolCallID,
				runtimeInputID:      ev.RuntimeInputID,
				stopReason:          ev.StopReason,
				model:               ev.Model,
				invocationID:        ev.ModelInvocationID,
				inputTokens:         ev.InputTokens,
				outputTokens:        ev.OutputTokens,
				cachedInputTokens:   ev.CachedInputTokens,
				activeContextTokens: ev.ActiveContextTokens,
			}
			if ev.ToolInput != nil {
				if s, ok := ev.ToolInput.(string); ok {
					collected.toolInput = s
				}
			}
			ch <- collected
		}
	}()
	return ch
}

func countKind(events []piCollectedEvent, kind string) int {
	n := 0
	for _, e := range events {
		if e.kind == kind {
			n++
		}
	}
	return n
}

func piTestConfig(t *testing.T, binary string) drivercontract.LaunchConfig {
	cfg := drivercontract.LaunchConfig{
		AgentID:         "test-agent",
		RuntimeType:     runtimeconfig.TypePi,
		Binary:          binary,
		WorkDir:         t.TempDir(),
		RuntimeStateDir: t.TempDir(),
		PermissionMode:  "full_access",
		Env:             map[string]string{},
	}
	return testLaunchConfig(cfg, testLaunchSurfaces{})
}

// TestPiRPCSingleTurnProjectsOnce verifies a full single assistant turn:
// user message surfacing, streaming text exactly once (delta + text_end +
// message_end reconciliation), one Usage, and exactly one root TurnEnd after
// agent_settled.
func TestPiRPCSingleTurnProjectsOnce(t *testing.T) {
	marker := "<!-- agentteam-runtime-input:msg-abc -->"
	scenario := []piFakeScript{
		piFakePromptScript(
			piFakeUserMessage("msg-abc", "hello\n\n"+marker),
			piFakeUserMessageEnd("msg-abc", "hello\n\n"+marker),
			piFakeAgentStart(),
			piFakeTurnStart(),
			piFakeAssistantStart(nil, nil),
			piFakeTextDelta(0, "Hello", piFakeUsage(10, 5, 0, 15)),
			piFakeTextDelta(0, " world", piFakeUsage(10, 5, 0, 15)),
			piFakeTextEnd(0, "Hello world", piFakeUsage(10, 5, 0, 15)),
			piFakeAssistantEnd([]map[string]any{piFakeTextBlock("Hello world")}, piFakeUsage(10, 5, 0, 15)),
			piFakeTurnEnd([]map[string]any{piFakeTextBlock("Hello world")}, piFakeUsage(10, 5, 0, 15)),
			piFakeAgentEnd(false),
			piFakeAgentSettled(),
		),
	}
	binary, _ := writeFakePiRPCModeBinary(t, scenario)
	d := &PiRPCDriver{}
	ctx := context.Background()
	cfg := piTestConfig(t, binary)
	result, err := d.Spawn(ctx, cfg)
	require.NoError(t, err)
	defer result.Cleanup()

	err = d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "msg-abc"}, "hello"), result.SessionID)
	require.NoError(t, err)

	events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle)
	require.Contains(t, []string{consts.ParsedEventKindTurnEnd, consts.ParsedEventKindRuntimeIdle}, events[len(events)-1].kind)

	// Exactly one root TurnEnd, and it must come after the text.
	require.Equal(t, 1, countKind(events, consts.ParsedEventKindTurnEnd), "exactly one AgentTeam root TurnEnd")
	// RuntimeIdle must follow the TurnEnd.
	turnIdx := -1
	idleIdx := -1
	for i, e := range events {
		if e.kind == consts.ParsedEventKindTurnEnd && turnIdx == -1 {
			turnIdx = i
		}
		if e.kind == consts.ParsedEventKindRuntimeIdle && idleIdx == -1 {
			idleIdx = i
		}
	}
	require.Greater(t, idleIdx, turnIdx, "RuntimeIdle after TurnEnd")

	// Streaming text appears exactly once with the full content.
	var textEvents []string
	for _, e := range events {
		if e.kind == consts.ParsedEventKindText {
			textEvents = append(textEvents, e.text)
		}
	}
	require.Equal(t, []string{"Hello", " world"}, textEvents, "append-only deltas once")
	require.NotEmpty(t, countKind(events, consts.ParsedEventKindText))

	// Surfacing correlated input once.
	acked := 0
	for _, e := range events {
		if e.kind == consts.ParsedEventKindRuntimeInputAcked && e.runtimeInputID == "msg-abc" {
			acked++
		}
	}
	require.GreaterOrEqual(t, acked, 1, "user message acked")

	// Exactly one Usage with a stable invocation identity.
	require.Equal(t, 1, countKind(events, consts.ParsedEventKindUsage))
	for _, e := range events {
		if e.kind == consts.ParsedEventKindUsage {
			require.Equal(t, uint64(15), e.activeContextTokens)
			require.Equal(t, "gpt-5.5", e.model)
			require.NotEmpty(t, e.invocationID, "stable ModelInvocationID")
		}
	}

	// No early terminal before agent_settled: the TurnEnd index must be after
	// the last tool/text activity.
	require.Greater(t, turnIdx, 0)
}

// TestPiRPCThinkingDeltasAggregateAsOneActivity verifies that token-sized Pi
// reasoning deltas are coalesced at the authoritative thinking_end boundary.
// Emitting every delta separately floods AgentTeam activity/trajectory reporting.
func TestPiRPCThinkingDeltasAggregateAsOneActivity(t *testing.T) {
	usage := piFakeUsage(10, 4, 0, 14)
	thinking := "Inspecting the request"
	scenario := []piFakeScript{
		piFakePromptScript(
			piFakeAgentStart(),
			piFakeTurnStart(),
			piFakeAssistantStart(nil, nil),
			piFakeThinkingDelta(0, "Inspecting", usage),
			piFakeThinkingDelta(0, " the", usage),
			piFakeThinkingDelta(0, " request", usage),
			piFakeThinkingEnd(0, thinking, usage),
			piFakeTextDelta(1, "done", usage),
			piFakeAssistantEnd([]map[string]any{piFakeThinkingBlock(thinking), piFakeTextBlock("done")}, usage),
			piFakeTurnEnd([]map[string]any{piFakeThinkingBlock(thinking), piFakeTextBlock("done")}, usage),
			piFakeAgentEnd(false),
			piFakeAgentSettled(),
		),
	}
	binary, _ := writeFakePiRPCModeBinary(t, scenario)
	d := &PiRPCDriver{}
	result, err := d.Spawn(context.Background(), piTestConfig(t, binary))
	require.NoError(t, err)
	defer result.Cleanup()

	require.NoError(t, d.DeliverRuntimeInput(context.Background(), testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "go"), result.SessionID))
	events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle)

	var got []string
	for _, event := range events {
		if event.kind == consts.ParsedEventKindThinking && event.text != "" {
			got = append(got, event.text)
		}
	}
	require.Equal(t, []string{thinking}, got)
}

// TestPiRPCMultiTurnTwoTools verifies an agent run with two Pi turn cycles and
// two tools. Each tool's id/input/output appears exactly once and no early root
// TurnEnd is emitted before agent_settled.
func TestPiRPCMultiTurnTwoTools(t *testing.T) {
	usage := piFakeUsage(20, 10, 0, 30)
	scenario := []piFakeScript{
		piFakePromptScript(
			piFakeAgentStart(),
			piFakeTurnStart(),
			piFakeAssistantStart([]map[string]any{piFakeToolCallBlock("call_a", "bash", map[string]any{"command": "ls"})}, usage),
			piFakeToolCallEnd(0, "call_a", "bash", map[string]any{"command": "ls"}, usage),
			piFakeToolExecutionStart("call_a", "bash", map[string]any{"command": "ls"}),
			piFakeToolExecutionUpdate("call_a", "bash", map[string]any{"command": "ls"}, map[string]any{"content": []any{map[string]any{"type": "text", "text": "a"}}}),
			piFakeToolExecutionUpdate("call_a", "bash", map[string]any{"command": "ls"}, map[string]any{"content": []any{map[string]any{"type": "text", "text": "ab"}}}),
			piFakeToolExecutionEnd("call_a", "bash", map[string]any{"content": []any{map[string]any{"type": "text", "text": "ab"}}}, false),
			piFakeAssistantEnd([]map[string]any{piFakeToolCallBlock("call_a", "bash", map[string]any{"command": "ls"})}, usage),
			piFakeTurnEnd([]map[string]any{piFakeToolCallBlock("call_a", "bash", map[string]any{"command": "ls"})}, usage),
			piFakeTurnStart(),
			piFakeAssistantStart([]map[string]any{piFakeTextBlock("done "), piFakeToolCallBlock("call_b", "read", map[string]any{"path": "x"})}, usage),
			piFakeTextDelta(0, "done ", usage),
			piFakeToolCallEnd(1, "call_b", "read", map[string]any{"path": "x"}, usage),
			piFakeToolExecutionStart("call_b", "read", map[string]any{"path": "x"}),
			piFakeToolExecutionEnd("call_b", "read", map[string]any{"content": []any{map[string]any{"type": "text", "text": "content of x"}}}, false),
			piFakeAssistantEnd([]map[string]any{piFakeTextBlock("done "), piFakeToolCallBlock("call_b", "read", map[string]any{"path": "x"})}, usage),
			piFakeTurnEnd([]map[string]any{piFakeTextBlock("done "), piFakeToolCallBlock("call_b", "read", map[string]any{"path": "x"})}, usage),
			piFakeAgentEnd(false),
			piFakeAgentSettled(),
		),
	}
	binary, _ := writeFakePiRPCModeBinary(t, scenario)
	d := &PiRPCDriver{}
	ctx := context.Background()
	cfg := piTestConfig(t, binary)
	result, err := d.Spawn(ctx, cfg)
	require.NoError(t, err)
	defer result.Cleanup()

	err = d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "run tools"), result.SessionID)
	require.NoError(t, err)

	events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle)

	// Exactly one root TurnEnd.
	require.Equal(t, 1, countKind(events, consts.ParsedEventKindTurnEnd))

	// Each tool id appears exactly once for tool_call and tool_output.
	var callA, callB, outA, outB int
	for _, e := range events {
		switch {
		case e.kind == consts.ParsedEventKindToolCall && e.toolCallID == "call_a":
			callA++
		case e.kind == consts.ParsedEventKindToolCall && e.toolCallID == "call_b":
			callB++
		case e.kind == consts.ParsedEventKindToolOutput && e.toolCallID == "call_a":
			outA++
		case e.kind == consts.ParsedEventKindToolOutput && e.toolCallID == "call_b":
			outB++
		}
	}
	require.Equal(t, 1, callA, "tool A call once")
	require.Equal(t, 1, callB, "tool B call once")
	require.Equal(t, 1, outA, "tool A output once")
	require.Equal(t, 1, outB, "tool B output once")

	// Tool output text is the final cumulative result, not partials.
	for _, e := range events {
		if e.kind == consts.ParsedEventKindToolOutput && e.toolCallID == "call_a" {
			require.Contains(t, e.text, "ab")
		}
	}
	// Text emitted once.
	require.Equal(t, 1, countKind(events, consts.ParsedEventKindText))
}

// TestPiRPCRetrySingleRootTurnEnd verifies an agent_end(willRetry=true)
// followed by auto_retry + a fresh agent run does NOT produce an early root
// TurnEnd; only the final agent_settled emits exactly one.
func TestPiRPCRetrySingleRootTurnEnd(t *testing.T) {
	usage := piFakeUsage(10, 5, 0, 15)
	scenario := []piFakeScript{
		piFakePromptScript(
			piFakeAgentStart(),
			piFakeTurnStart(),
			piFakeAssistantStart(nil, nil),
			piFakeTextDelta(0, "retry me", usage),
			piFakeAssistantEnd([]map[string]any{piFakeTextBlock("retry me")}, usage),
			piFakeTurnEnd([]map[string]any{piFakeTextBlock("retry me")}, usage),
			piFakeAgentEnd(true, piFakeAssistantMessage([]map[string]any{piFakeTextBlock("retry me")}, usage)),
			piFakeAutoRetryStart(1, 3),
			piFakeAgentStart(),
			piFakeTurnStart(),
			piFakeAssistantStart(nil, nil),
			piFakeTextDelta(0, "retried ok", usage),
			piFakeAssistantEnd([]map[string]any{piFakeTextBlock("retried ok")}, usage),
			piFakeTurnEnd([]map[string]any{piFakeTextBlock("retried ok")}, usage),
			piFakeAutoRetryEnd(true, 1),
			piFakeAgentEnd(false),
			piFakeAgentSettled(),
		),
	}
	binary, _ := writeFakePiRPCModeBinary(t, scenario)
	d := &PiRPCDriver{}
	ctx := context.Background()
	cfg := piTestConfig(t, binary)
	result, err := d.Spawn(ctx, cfg)
	require.NoError(t, err)
	defer result.Cleanup()

	err = d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "retry"), result.SessionID)
	require.NoError(t, err)

	events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle)
	require.Equal(t, 1, countKind(events, consts.ParsedEventKindTurnEnd), "exactly one root TurnEnd across retry")
	// The TurnEnd must appear after both text outputs.
	textIdx := 0
	turnIdx := 0
	for i, e := range events {
		if e.kind == consts.ParsedEventKindText {
			textIdx = i
		}
		if e.kind == consts.ParsedEventKindTurnEnd {
			turnIdx = i
		}
	}
	require.Greater(t, turnIdx, textIdx, "root TurnEnd after all retry work")
	// Both text deltas present (from both runs).
	require.Equal(t, 2, countKind(events, consts.ParsedEventKindText))
}

// TestPiRPCCompactionAndQueueProjected verifies compaction + queued steer do not
// emit an early root TurnEnd and the final assistant settles once.
func TestPiRPCCompactionAndQueueProjected(t *testing.T) {
	usage := piFakeUsage(10, 5, 0, 15)
	scenario := []piFakeScript{
		piFakePromptScript(
			piFakeAgentStart(),
			piFakeTurnStart(),
			piFakeAssistantStart(nil, nil),
			piFakeTextDelta(0, "before compact", usage),
			piFakeAssistantEnd([]map[string]any{piFakeTextBlock("before compact")}, usage),
			piFakeTurnEnd([]map[string]any{piFakeTextBlock("before compact")}, usage),
			piFakeAgentEnd(true),
			piFakeCompactionStart("overflow"),
			piFakeCompactionEnd("overflow"),
			piFakeAgentStart(),
			piFakeTurnStart(),
			piFakeAssistantStart(nil, nil),
			piFakeTextDelta(0, "after compact", usage),
			piFakeAssistantEnd([]map[string]any{piFakeTextBlock("after compact")}, usage),
			piFakeTurnEnd([]map[string]any{piFakeTextBlock("after compact")}, usage),
			piFakeAgentEnd(false),
			piFakeQueueUpdate([]string{"pending steer"}, nil),
			piFakeAgentSettled(),
		),
	}
	binary, _ := writeFakePiRPCModeBinary(t, scenario)
	d := &PiRPCDriver{}
	ctx := context.Background()
	cfg := piTestConfig(t, binary)
	result, err := d.Spawn(ctx, cfg)
	require.NoError(t, err)
	defer result.Cleanup()

	err = d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "compact"), result.SessionID)
	require.NoError(t, err)

	events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle)
	require.Equal(t, 1, countKind(events, consts.ParsedEventKindTurnEnd))
	require.Equal(t, 1, countKind(events, consts.ParsedEventKindRuntimeCompacting))
	require.Equal(t, 2, countKind(events, consts.ParsedEventKindText))
}

// TestPiRPCUsageCardinality verifies repeated message_update usage snapshots
// produce exactly one Usage per finalized assistant message with a stable
// identity.
func TestPiRPCUsageCardinality(t *testing.T) {
	scenario := []piFakeScript{
		piFakePromptScript(
			piFakeAgentStart(),
			piFakeTurnStart(),
			piFakeAssistantStart(nil, nil),
			piFakeTextDelta(0, "a", piFakeUsage(1, 1, 0, 2)),
			piFakeTextDelta(0, "b", piFakeUsage(2, 2, 0, 4)),
			piFakeTextDelta(0, "c", piFakeUsage(3, 3, 0, 6)),
			piFakeAssistantEnd([]map[string]any{piFakeTextBlock("abc")}, piFakeUsage(3, 3, 0, 6)),
			piFakeTurnEnd([]map[string]any{piFakeTextBlock("abc")}, piFakeUsage(3, 3, 0, 6)),
			piFakeAgentEnd(false),
			piFakeAgentSettled(),
		),
	}
	binary, _ := writeFakePiRPCModeBinary(t, scenario)
	d := &PiRPCDriver{}
	ctx := context.Background()
	cfg := piTestConfig(t, binary)
	result, err := d.Spawn(ctx, cfg)
	require.NoError(t, err)
	defer result.Cleanup()

	err = d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "usage"), result.SessionID)
	require.NoError(t, err)

	events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle)
	require.Equal(t, 1, countKind(events, consts.ParsedEventKindUsage), "one Usage per finalized assistant message")
	for _, e := range events {
		if e.kind == consts.ParsedEventKindUsage {
			require.Equal(t, uint64(3), e.inputTokens, "latest cumulative input")
			require.Equal(t, uint64(6), e.activeContextTokens, "latest cumulative total")
			require.NotEmpty(t, e.invocationID, "stable identity")
		}
	}
}

// TestPiRPCPromptAcceptanceBeforeAgentStart verifies DeliverRuntimeInput
// returns on the prompt preflight response (before agent_start), and that an
// abort during that gap does not falsely observe idle.
func TestPiRPCPromptAcceptanceBeforeAgentStart(t *testing.T) {
	scenario := []piFakeScript{
		piFakePromptScript(), // response only, no events
		piFakeAbortScript(piFakeAgentEnd(false), piFakeAgentSettled()),
		piFakePromptScript(
			piFakeAgentStart(),
			piFakeTurnStart(),
			piFakeAssistantStart(nil, nil),
			piFakeTextDelta(0, "ok", piFakeUsage(1, 1, 0, 2)),
			piFakeAssistantEnd([]map[string]any{piFakeTextBlock("ok")}, piFakeUsage(1, 1, 0, 2)),
			piFakeTurnEnd([]map[string]any{piFakeTextBlock("ok")}, piFakeUsage(1, 1, 0, 2)),
			piFakeAgentEnd(false),
			piFakeAgentSettled(),
		),
	}
	binary, _ := writeFakePiRPCModeBinary(t, scenario)
	d := &PiRPCDriver{}
	ctx := context.Background()
	cfg := piTestConfig(t, binary)
	result, err := d.Spawn(ctx, cfg)
	require.NoError(t, err)
	defer result.Cleanup()

	// First prompt: acceptance returns before any agent_start event.
	err = d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "first"), result.SessionID)
	require.NoError(t, err)
	// Root run is active (agent events have not arrived yet).
	d.mu.Lock()
	active := d.rootRunActive
	d.mu.Unlock()
	require.True(t, active, "root run active after prompt acceptance")

	// Abort during the gap waits for the same generation's agent_settled.
	abortCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()
	err = d.InterruptActiveRun(abortCtx, result.SessionID)
	require.NoError(t, err)

	// After abort settles, a second prompt starts a fresh generation.
	err = d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "second"), result.SessionID)
	require.NoError(t, err)
	events := collectPiEventsN(t, result, consts.ParsedEventKindRuntimeIdle, 2)
	require.Equal(t, 2, countKind(events, consts.ParsedEventKindTurnEnd), "two root runs produce two root TurnEnds")
}

// TestPiRPCDeliverAckedBusySteerAcknowledges verifies a busy steer is acknowledged
// with a acked event for the correlated id.
func TestPiRPCDeliverAckedBusySteerAcknowledges(t *testing.T) {
	marker := "<!-- agentteam-runtime-input:steer-1 -->"
	scenario := []piFakeScript{
		piFakePromptScript(
			piFakeAgentStart(),
			piFakeTurnStart(),
			piFakeAssistantStart(nil, nil),
			piFakeTextDelta(0, "busy", piFakeUsage(1, 1, 0, 2)),
		),
		piFakeSteerScript(
			piFakeUserMessage("steer-1", "steer\n\n"+marker),
			piFakeUserMessageEnd("steer-1", "steer\n\n"+marker),
			piFakeAssistantEnd([]map[string]any{piFakeTextBlock("busy")}, piFakeUsage(1, 1, 0, 2)),
			piFakeTurnEnd([]map[string]any{piFakeTextBlock("busy")}, piFakeUsage(1, 1, 0, 2)),
			piFakeAgentEnd(false),
			piFakeAgentSettled(),
		),
	}
	binary, _ := writeFakePiRPCModeBinary(t, scenario)
	d := &PiRPCDriver{}
	ctx := context.Background()
	cfg := piTestConfig(t, binary)
	result, err := d.Spawn(ctx, cfg)
	require.NoError(t, err)
	defer result.Cleanup()

	// Start a root run and let it become active.
	err = d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "start"), result.SessionID)
	require.NoError(t, err)

	// Busy steer while the root run is active; the method returns as soon as
	// the steer RPC is accepted (not when the marker is acknowledged).
	steerDone := make(chan error, 1)
	go func() {
		steerDone <- d.DeliverAckedBusyRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "steer-1"}, "steer"), result.SessionID)
	}()
	select {
	case err := <-steerDone:
		require.NoError(t, err, "steer returns on RPC acceptance")
	case <-time.After(5 * time.Second):
		t.Fatal("DeliverAckedBusyRuntimeInput did not return on steer acceptance")
	}
	// Collect events; the steer acked event must carry steer-1.
	events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle)
	found := false
	for _, e := range events {
		if e.kind == consts.ParsedEventKindRuntimeInputAcked && e.runtimeInputID == "steer-1" {
			found = true
		}
	}
	require.True(t, found, "steer acked for correlated id")
	require.Equal(t, 1, countKind(events, consts.ParsedEventKindTurnEnd))
}

// TestPiRPCRootRunRemainsBusyUntilSettled verifies the driver does not emit a
// TurnEnd on a Pi turn_end or agent_end while a root run is active.
func TestPiRPCRootRunRemainsBusyUntilSettled(t *testing.T) {
	usage := piFakeUsage(1, 1, 0, 2)
	scenario := []piFakeScript{
		piFakePromptScript(
			piFakeAgentStart(),
			piFakeTurnStart(),
			piFakeAssistantStart(nil, nil),
			piFakeTextDelta(0, "x", usage),
			piFakeAssistantEnd([]map[string]any{piFakeTextBlock("x")}, usage),
			piFakeTurnEnd([]map[string]any{piFakeTextBlock("x")}, usage),
			piFakeAgentEnd(false),
			// No agent_settled yet: root stays active.
		),
	}
	binary, _ := writeFakePiRPCModeBinary(t, scenario)
	d := &PiRPCDriver{}
	ctx := context.Background()
	cfg := piTestConfig(t, binary)
	result, err := d.Spawn(ctx, cfg)
	require.NoError(t, err)
	defer result.Cleanup()

	err = d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "x"), result.SessionID)
	require.NoError(t, err)

	// Collect until the stream goes quiet (the scenario emits turn_end and
	// agent_end but no agent_settled). If the driver wrongly projected an early
	// terminal, a TurnEnd would appear in the collected events.
	events := collectPiEventsUntilQuiet(t, result)
	require.Equal(t, 1, countKind(events, consts.ParsedEventKindText), "text emitted")
	require.Equal(t, 0, countKind(events, consts.ParsedEventKindTurnEnd), "no TurnEnd before agent_settled")

	d.mu.Lock()
	active := d.rootRunActive
	d.mu.Unlock()
	require.True(t, active, "root run still active without agent_settled")
}

// collectPiEventsUntilQuiet reads translated events until a quiet window with no
// new event (or the deadline) passes. The scenario for the busy test emits no
// agent_settled, so after the events settle the stream goes quiet.
func collectPiEventsUntilQuiet(t *testing.T, result drivercontract.SpawnResult) []piCollectedEvent {
	t.Helper()
	var events []piCollectedEvent
	deadline := time.After(10 * time.Second)
	stream := piStreamEvents(t, result)
	for {
		select {
		case ev, ok := <-stream:
			if !ok {
				return events
			}
			events = append(events, ev)
		case <-time.After(500 * time.Millisecond):
			return events
		case <-deadline:
			return events
		}
	}
}

// TestPiRPCStaleTerminalIgnored verifies a late turn_end/agent_end after
// agent_settled is ignored (no duplicate root TurnEnd).
func TestPiRPCStaleTerminalIgnored(t *testing.T) {
	usage := piFakeUsage(1, 1, 0, 2)
	scenario := []piFakeScript{
		piFakePromptScript(
			piFakeAgentStart(),
			piFakeTurnStart(),
			piFakeAssistantStart(nil, nil),
			piFakeTextDelta(0, "ok", usage),
			piFakeAssistantEnd([]map[string]any{piFakeTextBlock("ok")}, usage),
			piFakeTurnEnd([]map[string]any{piFakeTextBlock("ok")}, usage),
			piFakeAgentEnd(false),
			piFakeAgentSettled(),
			// Stale events after settlement.
			piFakeTurnStart(),
			piFakeTurnEnd([]map[string]any{piFakeTextBlock("stale")}, usage),
			piFakeAgentEnd(false),
		),
	}
	binary, _ := writeFakePiRPCModeBinary(t, scenario)
	d := &PiRPCDriver{}
	ctx := context.Background()
	cfg := piTestConfig(t, binary)
	result, err := d.Spawn(ctx, cfg)
	require.NoError(t, err)
	defer result.Cleanup()

	err = d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "x"), result.SessionID)
	require.NoError(t, err)
	events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle)
	require.Equal(t, 1, countKind(events, consts.ParsedEventKindTurnEnd), "stale terminal events ignored")
	require.Equal(t, 1, countKind(events, consts.ParsedEventKindRuntimeIdle))
}

// TestPiRPCTransportFailureBeforeAcceptance verifies a transport failure while a
// prompt is pending resolves the delivery with an error (no root run started).
func TestPiRPCTransportFailureBeforeAcceptance(t *testing.T) {
	scenario := []piFakeScript{
		piFakeScript{Cmd: "prompt"}, // no response: command stays pending
	}
	binary, _ := writeFakePiRPCModeBinary(t, scenario)
	d := &PiRPCDriver{}
	ctx := context.Background()
	cfg := piTestConfig(t, binary)
	result, err := d.Spawn(ctx, cfg)
	require.NoError(t, err)

	deliverDone := make(chan error, 1)
	go func() {
		deliverDone <- d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "x"), result.SessionID)
	}()
	// Wait until the prompt command is pending (accepted registration visible),
	// then tear down the transport.
	require.Eventually(t, func() bool {
		d.mu.Lock()
		pending := d.promptAccepted != nil
		d.mu.Unlock()
		return pending
	}, 5*time.Second, 20*time.Millisecond)
	result.Terminate()
	deliverErr := <-deliverDone
	require.Error(t, deliverErr, "pending delivery resolves with an error")
}

// TestPiRPCTransportFailureMidRun verifies a transport failure after prompt
// acceptance (root run active, no agent_settled) emits Error + terminal TurnEnd
// + RuntimeIdle.
func TestPiRPCTransportFailureMidRun(t *testing.T) {
	scenario := []piFakeScript{
		piFakePromptScript(piFakeAgentStart(), piFakeTurnStart()),
	}
	binary, _ := writeFakePiRPCModeBinary(t, scenario)
	d := &PiRPCDriver{}
	ctx := context.Background()
	cfg := piTestConfig(t, binary)
	result, err := d.Spawn(ctx, cfg)
	require.NoError(t, err)

	err = d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "x"), result.SessionID)
	require.NoError(t, err)
	// Mid-run transport teardown (no agent_settled).
	result.Terminate()

	events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle)
	require.Equal(t, 1, countKind(events, consts.ParsedEventKindTurnEnd), "transport failure emits terminal TurnEnd")
	require.Equal(t, 1, countKind(events, consts.ParsedEventKindRuntimeIdle))
	require.Equal(t, 1, countKind(events, consts.ParsedEventKindError))
	_ = io.ErrClosedPipe
}

// TestPiRPCParseLineAckedRequiresInputID
func TestPiRPCParseLineAckedRequiresInputID(t *testing.T) {
	d := &PiRPCDriver{}
	require.Empty(t, d.ParseLine(`{"kind":"runtime_input_surfaced","sessionId":"s1"}`))
}