package agent

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

type idleCompactionProcessStubDriver struct {
	started  bool
	startErr error
	calls    int
	sessions []string
}

func (d *idleCompactionProcessStubDriver) Name() string { return "codex" }

func (d *idleCompactionProcessStubDriver) SupportsStdinNotification() bool { return false }

func (d *idleCompactionProcessStubDriver) MCPToolPrefix() string { return "mcp__agentteam__" }

func (d *idleCompactionProcessStubDriver) Spawn(context.Context, drivercontract.LaunchConfig) (drivercontract.SpawnResult, error) {
	return drivercontract.SpawnResult{}, nil
}

func (d *idleCompactionProcessStubDriver) ParseLine(string) []drivercontract.ParsedEvent { return nil }

func (d *idleCompactionProcessStubDriver) EncodeStdinMessage(string, string) []byte { return nil }

func (d *idleCompactionProcessStubDriver) StartIdleCompaction(_ context.Context, sessionID string) (bool, error) {
	d.calls++
	d.sessions = append(d.sessions, sessionID)
	return d.started, d.startErr
}

func TestHandleParsedEventTurnEndEmitsRuntimeCompactingWhenIdleCompactionStarts(t *testing.T) {
	driver := &idleCompactionProcessStubDriver{started: true}

	var events []drivercontract.ParsedEvent
	ap := &Process{
		AgentID:   "agent-a",
		SessionID: "sess-1",
		processRuntimeBinding: processRuntimeBinding{driver: driver,
			onEvent: func(ev drivercontract.ParsedEvent) {
				events = append(events, ev)
			}}}
	require.NoError(t, ap.SetOnResumeSafety(func(uint64, bool, string) error { return nil }))

	ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})

	require.Equal(t, 1, driver.calls)
	require.Equal(t, []string{"sess-1"}, driver.sessions)
	require.False(t, ap.isIdle)
	require.Len(t, events, 1)
	require.Equal(t, consts.ParsedEventKindRuntimeCompacting, events[0].Kind)
	require.Equal(t, "sess-1", events[0].SessionID)
}

func TestHandleParsedEventTurnEndDoesNotCompactAheadOfPendingRuntimeInput(t *testing.T) {
	driver := &idleCompactionProcessStubDriver{started: true}
	ap := &Process{
		AgentID:               "agent-pending-before-compact",
		SessionID:             "sess-pending-before-compact",
		processRuntimeBinding: processRuntimeBinding{driver: driver},
	}
	ap.pendingRuntimeInputs = map[string]pendingRuntimeInput{
		"message-1": {
			Input:          testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "message-1"}, "accepted input"),
			InflightReason: DeliveryInflightOutcomeUnknown,
		},
	}
	require.NoError(t, ap.SetOnResumeSafety(func(uint64, bool, string) error { return nil }))

	ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})

	require.Zero(t, driver.calls,
		"a new compaction mutation must remain behind the older pending ownership record")
}

// proactiveCompactionStubDriver implements only
// drivercontract.ProactiveIdleCompactionDriver. It intentionally does not
// implement IdleCompactionDriver so the overflow path is bypassed and the
// proactive between-turn path is exercised directly. Its Name is non-codex so
// resume-safety persistence is a no-op and no persister callback is required.
type proactiveCompactionStubDriver struct {
	started  bool
	startErr error
	onStart  func()
	calls    int
	sessions []string
}

func (d *proactiveCompactionStubDriver) Name() string                    { return "opencode" }
func (d *proactiveCompactionStubDriver) SupportsStdinNotification() bool { return false }
func (d *proactiveCompactionStubDriver) MCPToolPrefix() string           { return "mcp_agentteam_" }
func (d *proactiveCompactionStubDriver) Spawn(context.Context, drivercontract.LaunchConfig) (drivercontract.SpawnResult, error) {
	return drivercontract.SpawnResult{}, nil
}
func (d *proactiveCompactionStubDriver) ParseLine(string) []drivercontract.ParsedEvent { return nil }
func (d *proactiveCompactionStubDriver) EncodeStdinMessage(string, string) []byte      { return nil }

func (d *proactiveCompactionStubDriver) StartProactiveIdleCompaction(_ context.Context, sessionID string) (bool, error) {
	d.calls++
	d.sessions = append(d.sessions, sessionID)
	if d.onStart != nil {
		d.onStart()
	}
	return d.started, d.startErr
}

type proactiveCompactionDeliveryStubDriver struct {
	proactiveCompactionStubDriver
	inputs           []drivercontract.RuntimeInput
	deliverySessions []string
}

func (d *proactiveCompactionDeliveryStubDriver) SupportsStdinNotification() bool { return true }

func (d *proactiveCompactionDeliveryStubDriver) DeliverBusyStdinInput(
	context.Context, drivercontract.RuntimeInput, string,
) error {
	return nil
}

func (d *proactiveCompactionDeliveryStubDriver) DeliverRuntimeInput(
	_ context.Context,
	input drivercontract.RuntimeInput,
	sessionID string,
) error {
	d.inputs = append(d.inputs, input)
	d.deliverySessions = append(d.deliverySessions, sessionID)
	return nil
}

func TestHandleParsedEventTurnEndChecksProactiveCompactionAfterConfirmedWakeReply(t *testing.T) {
	driver := &proactiveCompactionStubDriver{started: true}
	ap := newTestProcessForRecovery("agent-proactive-wake", "sess-1", driver)
	ap.mu.Lock()
	_, started := ap.beginWakeDeliveryForMessageLocked(types.InboxMessage{
		MessageID: "msg-1",
		Target:    "#agentteam:00000014",
		TargetKey: "ops",
		TargetSeq: 1,
	}, "received")
	ap.mu.Unlock()
	require.True(t, started)
	require.True(t, ap.ConfirmVisibleReplyForTarget("ops", 1, "reply-1"))

	ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})

	require.Equal(t, 1, driver.calls)
	require.Equal(t, []string{"sess-1"}, driver.sessions)
}

func TestHandleParsedEventTurnEndRunsFinalCheckBeforeProactiveCompaction(t *testing.T) {
	driver := &proactiveCompactionDeliveryStubDriver{
		proactiveCompactionStubDriver: proactiveCompactionStubDriver{started: true},
	}
	ap := newTestProcessForRecovery("agent-proactive-final-check", "sess-1", driver)
	ap.mu.Lock()
	_, started := ap.beginWakeDeliveryForMessageLocked(types.InboxMessage{
		MessageID: "msg-1",
		Target:    "#agentteam:00000014",
	}, "received")
	ap.mu.Unlock()
	require.True(t, started)

	ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})

	require.Len(t, driver.inputs, 1)
	require.Zero(t, driver.calls)

	ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})

	require.Equal(t, 1, driver.calls)
}

func TestHandleParsedEventTurnEndEmitsRuntimeCompactingWhenProactiveCompactionStarts(t *testing.T) {
	driver := &proactiveCompactionStubDriver{started: true}

	var events []drivercontract.ParsedEvent
	ap := newTestProcessForRecovery("agent-proactive", "sess-1", driver)
	ap.processRuntimeBinding.onEvent = func(ev drivercontract.ParsedEvent) {
		events = append(events, ev)
	}

	ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})

	require.Equal(t, 1, driver.calls)
	require.Equal(t, []string{"sess-1"}, driver.sessions)
	require.False(t, ap.isIdle)
	require.Contains(t, parsedEventKinds(events), consts.ParsedEventKindRuntimeCompacting)
}

func TestHandleParsedEventTurnEndTreatsAmbiguousProactiveCompactionAsActive(t *testing.T) {
	driver := &proactiveCompactionStubDriver{
		started:  true,
		startErr: context.DeadlineExceeded,
	}
	var events []drivercontract.ParsedEvent
	ap := newTestProcessForRecovery("agent-proactive-ambiguous", "sess-1", driver)
	ap.processRuntimeBinding.onEvent = func(ev drivercontract.ParsedEvent) {
		events = append(events, ev)
	}

	ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})

	require.Equal(t, 1, driver.calls)
	require.False(t, ap.isIdle)
	require.Contains(t, parsedEventKinds(events), consts.ParsedEventKindRuntimeCompacting)
}

func TestHandleParsedEventTurnEndReturnsIdleAfterRejectedProactiveCompaction(t *testing.T) {
	driver := &proactiveCompactionStubDriver{startErr: context.Canceled}
	var events []drivercontract.ParsedEvent
	ap := newTestProcessForRecovery("agent-proactive-rejected", "sess-1", driver)
	ap.processRuntimeBinding.onEvent = func(ev drivercontract.ParsedEvent) {
		events = append(events, ev)
	}

	ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})

	require.Equal(t, 1, driver.calls)
	require.True(t, ap.isIdle)
	require.NotContains(t, parsedEventKinds(events), consts.ParsedEventKindRuntimeCompacting)
}

func TestHandleParsedEventTurnEndProactiveCompactionDoesNotInjectContinuation(t *testing.T) {
	driver := &proactiveCompactionStubDriver{started: true}
	stdin := &recordingWriteCloser{}
	var events []drivercontract.ParsedEvent
	ap := newTestProcessForRecovery("agent-proactive-noretry", "sess-1", driver)
	ap.stdinPipe = stdin
	ap.processRuntimeBinding.onEvent = func(ev drivercontract.ParsedEvent) {
		events = append(events, ev)
	}

	// A plain turn_end (no overflow) triggers proactive between-turn compaction.
	ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})
	require.Contains(t, parsedEventKinds(events), consts.ParsedEventKindRuntimeCompacting)
	require.False(t, ap.isIdle)

	// Compaction completes via runtime_idle. Unlike the overflow path, no
	// continuation retry is injected: the process simply returns to idle and the
	// next user turn continues from the compacted context.
	ap.handleParsedEvent(drivercontract.ParsedEvent{
		Kind:      consts.ParsedEventKindRuntimeIdle,
		SessionID: "sess-1",
	})

	require.Empty(t, stdin.String(), "proactive between-turn compaction must not inject a continuation retry")
	require.True(t, ap.isIdle)
	require.Equal(t, 1, driver.calls, "compaction completion must not start another idle-compaction check")
}

func TestHandleParsedEventTurnEndSkipsProactiveCompactionWhenRecoveryDisabled(t *testing.T) {
	driver := &proactiveCompactionStubDriver{started: true}
	ap := newTestProcessForRecovery("agent-proactive-disabled", "sess-1", driver)
	ap.disableRecovery = true

	ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})

	require.Zero(t, driver.calls, "proactive compaction must be skipped when recovery is disabled")
}

func TestHandleParsedEventTurnEndDoesNotCompactProactivelyAheadOfPendingRuntimeInput(t *testing.T) {
	driver := &proactiveCompactionStubDriver{started: true}
	ap := newTestProcessForRecovery("agent-proactive-pending", "sess-1", driver)
	ap.pendingRuntimeInputs = map[string]pendingRuntimeInput{
		"message-1": {
			Input:          testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "message-1"}, "accepted input"),
			InflightReason: DeliveryInflightOutcomeUnknown,
		},
	}

	ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})

	require.Zero(t, driver.calls)
}

func TestDeliveryArrivingAtProactiveCompactionBoundaryWaitsAndRunsAfterCompletion(t *testing.T) {
	startEntered := make(chan struct{})
	startRelease := make(chan struct{})
	driver := &proactiveCompactionDeliveryStubDriver{
		proactiveCompactionStubDriver: proactiveCompactionStubDriver{
			started: true,
			onStart: func() {
				close(startEntered)
				<-startRelease
			},
		},
	}
	ap := newTestProcessForRecovery("agent-proactive-delivery", "sess-1", driver)
	ap.workDir = t.TempDir()

	turnEndDone := make(chan struct{})
	go func() {
		defer close(turnEndDone)
		ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})
	}()
	select {
	case <-startEntered:
	case <-time.After(time.Second):
		t.Fatal("proactive compaction did not reach the mutation boundary")
	}
	select {
	case <-ap.deliveryGate:
		ap.releaseDeliveryGate()
		close(startRelease)
		<-turnEndDone
		t.Fatal("proactive compaction did not retain the delivery admission gate")
	default:
	}

	deliveryDone := make(chan error, 1)
	go func() {
		deliveryDone <- ap.DeliverForHandoff(t.Context(), types.InboxMessage{
			MessageID:   "message-during-compaction",
			Target:      "#agentteam:00000014",
			ChannelType: "channel",
			ChannelName: "agentteam",
			SenderType:  "user",
			SenderName:  "owner",
			Content:     "new work",
			Timestamp:   time.Now().UTC().Format(time.RFC3339),
		})
	}()
	close(startRelease)
	select {
	case <-turnEndDone:
	case <-time.After(time.Second):
		t.Fatal("turn end did not complete after proactive compaction started")
	}
	select {
	case err := <-deliveryDone:
		require.NoError(t, err)
	case <-time.After(time.Second):
		t.Fatal("delivery did not resume after proactive compaction released the gate")
	}
	ap.mu.Lock()
	require.Len(t, ap.inbox, 1)
	ap.mu.Unlock()
	require.Empty(t, driver.inputs)

	ap.handleParsedEvent(drivercontract.ParsedEvent{
		Kind:      consts.ParsedEventKindRuntimeIdle,
		SessionID: "sess-1",
	})

	require.Len(t, driver.inputs, 1)
	require.Equal(t, "message-during-compaction", driver.inputs[0].ID)
	require.Equal(t, []string{"sess-1"}, driver.deliverySessions)
	require.Equal(t, "sess-1", ap.GetSessionID())
}

func TestProactiveCompactionFailureDoesNotTriggerContextRecoveryOrRecheck(t *testing.T) {
	driver := &proactiveCompactionDeliveryStubDriver{
		proactiveCompactionStubDriver: proactiveCompactionStubDriver{started: true},
	}
	ap := newTestProcessForRecovery("agent-proactive-failure", "sess-1", driver)
	ap.workDir = t.TempDir()
	var events []drivercontract.ParsedEvent
	ap.processRuntimeBinding.onEvent = func(ev drivercontract.ParsedEvent) {
		events = append(events, ev)
	}

	ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})
	require.Equal(t, 1, driver.calls)
	require.NotContains(t, parsedEventKinds(events), consts.ParsedEventKindTurnEnd)

	failureMetadata := map[string]any{
		contextCompactionFailedMetadataKey:      true,
		contextCompactionMaintenanceMetadataKey: true,
	}
	ap.handleParsedEvent(drivercontract.ParsedEvent{
		Kind:     consts.ParsedEventKindError,
		Metadata: failureMetadata,
	})
	ap.handleParsedEvent(drivercontract.ParsedEvent{
		Kind:     consts.ParsedEventKindTurnEnd,
		Metadata: failureMetadata,
	})

	require.False(t, ap.stopping)
	require.False(t, ap.sleeping)
	require.False(t, ap.contextWindowExceeded)
	require.False(t, ap.isIdle, "a retired runtime generation must stay non-idle until process exit")
	require.Equal(t, 1, driver.calls, "maintenance failure boundary must not recursively check compaction")
	require.NotContains(t, parsedEventKinds(events), consts.ParsedEventKindTurnEnd,
		"maintenance failure must not publish an externally idle turn boundary")

	require.NoError(t, ap.DeliverForHandoff(t.Context(), types.InboxMessage{
		MessageID:   "message-after-failed-compaction",
		Target:      "#agentteam:00000014",
		ChannelType: "channel",
		ChannelName: "agentteam",
		SenderType:  "user",
		SenderName:  "owner",
		Content:     "new work",
		Timestamp:   time.Now().UTC().Format(time.RFC3339),
	}))
	require.Empty(t, driver.inputs, "delivery must not mutate a retiring runtime generation")
	ap.mu.Lock()
	require.Len(t, ap.inbox, 1)
	ap.mu.Unlock()
}

func TestProactiveIdleCompactionCapabilityMapping(t *testing.T) {
	tests := []struct {
		name   string
		driver drivercontract.Driver
		want   bool
	}{
		{name: "OpenCode", driver: &OpenCodeDriver{}, want: true},
		{name: "Codex", driver: &CodexAppServerDriver{}, want: true},
		{name: "Claude", driver: &ClaudeDriver{}},
		{name: "Kimi", driver: &KimiWireDriver{}},
		{name: "Pi", driver: &PiRPCDriver{}, want: false},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			_, got := test.driver.(drivercontract.ProactiveIdleCompactionDriver)
			require.Equal(t, test.want, got)
		})
	}
}