package agent

import (
	"encoding/json"
	"testing"

	"github.com/stretchr/testify/require"

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

func TestPiRPCShortSessionCompactionNoopIsNonErrorActivity(t *testing.T) {
	scenario := []piFakeScript{piFakePromptScript(
		piFakeAgentStart(), piFakeTurnStart(),
		piFakeCompactionStart("manual"), piFakeCompactionEndTooSmall(),
		piFakeAgentEnd(false), piFakeAgentSettled(),
	)}
	binary, _ := writeFakePiRPCModeBinary(t, scenario)
	d := &PiRPCDriver{}
	result, err := d.Spawn(t.Context(), piTestConfig(t, binary))
	require.NoError(t, err)
	t.Cleanup(result.Cleanup)
	require.NoError(t, d.DeliverRuntimeInput(t.Context(), testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "compact"), result.SessionID))

	events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle)
	require.Zero(t, countKind(events, consts.ParsedEventKindError))
	require.Equal(t, 1, countKind(events, consts.ParsedEventKindTurnEnd))
	found := false
	for _, event := range events {
		if event.kind == consts.ParsedEventKindThinking && event.text == "Nothing to compact (session too small)" {
			found = true
		}
		require.NotContains(t, event.message, "Nothing to compact")
	}
	require.True(t, found)
}

func drainPiCompactionEvents(output *piTranslatedOutput) []translatedEvent {
	events := make([]translatedEvent, 0)
	for {
		select {
		case line := <-output.lines:
			var event translatedEvent
			if json.Unmarshal([]byte(line), &event) == nil {
				events = append(events, event)
			}
		default:
			return events
		}
	}
}

func TestPiRPCStandaloneNativeCompactionCompletionEmitsIdle(t *testing.T) {
	output := newPiTranslatedOutput(16)
	d := &PiRPCDriver{sessionID: "sess-native", output: output}
	d.onCompactionStart("sess-native", piRPCFrame{Reason: "manual"})
	d.onCompactionEnd("sess-native", piRPCFrame{
		Reason: "manual",
		Result: json.RawMessage(`{"summary":"compacted"}`),
	})

	events := drainPiCompactionEvents(output)
	require.Equal(t, 1, countTranslatedPiEvents(events, consts.ParsedEventKindRuntimeCompacting))
	require.Equal(t, 1, countTranslatedPiEvents(events, consts.ParsedEventKindRuntimeIdle))
	for _, event := range events {
		if event.Kind == consts.ParsedEventKindThinking {
			require.Equal(t, "Pi context compaction completed", event.Text)
			require.True(t, parsedBoolMetadata(event.Metadata, contextCompactionConfirmedMetadataKey))
		}
	}
}

func TestPiRPCCompactionCompletionDoesNotIdleActiveOrQueuedRoot(t *testing.T) {
	tests := []struct {
		name  string
		setup func(*PiRPCDriver)
	}{
		{
			name: "active root",
			setup: func(d *PiRPCDriver) {
				d.rootRunActive = true
				d.rootRunGeneration = 1
				d.rootRunSettled = make(chan struct{})
			},
		},
		{
			name: "queued delivery",
			setup: func(d *PiRPCDriver) {
				d.steerInFlight = map[string]bool{"queued": true}
			},
		},
		{
			name: "prompt preflight",
			setup: func(d *PiRPCDriver) {
				d.promptAccepted = make(chan error, 1)
			},
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			output := newPiTranslatedOutput(16)
			d := &PiRPCDriver{sessionID: "sess-native", output: output}
			test.setup(d)
			d.onCompactionStart("sess-native", piRPCFrame{Reason: "manual"})
			d.onCompactionEnd("sess-native", piRPCFrame{
				Reason: "manual",
				Result: json.RawMessage(`{"summary":"compacted"}`),
			})
			require.Zero(t, countTranslatedPiEvents(drainPiCompactionEvents(output), consts.ParsedEventKindRuntimeIdle))
		})
	}
}

func TestPiRPCNativeCompactionNoopEmitsStandaloneIdle(t *testing.T) {
	output := newPiTranslatedOutput(16)
	d := &PiRPCDriver{sessionID: "sess-native", output: output}
	d.onCompactionStart("sess-native", piRPCFrame{Reason: "manual"})
	d.onCompactionEnd("sess-native", piRPCFrame{
		Reason:       "manual",
		ErrorMessage: "Compaction failed: Nothing to compact (session too small)",
	})

	events := drainPiCompactionEvents(output)
	require.Equal(t, 1, countTranslatedPiEvents(events, consts.ParsedEventKindRuntimeIdle))
	found := false
	for _, event := range events {
		if event.Kind == consts.ParsedEventKindThinking && event.Text == "Nothing to compact (session too small)" {
			found = true
		}
	}
	require.True(t, found)
}

func countTranslatedPiEvents(events []translatedEvent, kind string) int {
	count := 0
	for _, event := range events {
		if event.Kind == kind {
			count++
		}
	}
	return count
}