package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"strings"
	"sync"
	"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"
)

func TestKimiACPProcessIdleDeliveryUsesPromptRPC(t *testing.T) {
	var stdin bytes.Buffer
	driver := newKimiACPTestDriver(&stdin)
	driver.sessionID = kimiTestSessionID
	store := &memoryDeliveryOwnershipStore{}
	process := newKimiProcessDeliveryTestProcess(t, driver, store, true, nil)

	_, hasProtocolDelivery := any(driver).(drivercontract.BusyStdinMessageDriver)
	require.True(t, hasProtocolDelivery, "idle Kimi ACP delivery must select session/prompt")
	_, supportsBusyDelivery := any(driver).(drivercontract.BusyStdinDeliveryDriver)
	require.False(t, supportsBusyDelivery, "Kimi ACP must queue input while a prompt is active")

	require.NoError(t, process.DeliverForHandoff(context.Background(), types.InboxMessage{
		MessageID: "kimi-acp-idle-1",
		Target:    "#agentteam:thread",
		Content:   "handle the ACP follow-up",
	}))

	var request kimiACPEnvelope
	require.NoError(t, json.Unmarshal(bytes.TrimSpace(stdin.Bytes()), &request))
	require.Equal(t, "session/prompt", request.Method)
	require.Contains(t, string(request.Params), "handle the ACP follow-up")
	require.NotEqual(t, "{}", strings.TrimSpace(stdin.String()))
	requireKimiOwnershipAcked(t, store, "kimi-acp-idle-1")

	driver.routeEnvelope(kimiACPEnvelope{
		ID: request.ID, Result: json.RawMessage(`{"stopReason":"end_turn"}`),
	})
}

type kimiProcessRPCRequest struct {
	method string
	params any
}

type kimiProcessEventRecorder struct {
	mu     sync.Mutex
	events []drivercontract.ParsedEvent
}

func (r *kimiProcessEventRecorder) record(event drivercontract.ParsedEvent) {
	r.mu.Lock()
	r.events = append(r.events, event)
	r.mu.Unlock()
}

func (r *kimiProcessEventRecorder) snapshot() []drivercontract.ParsedEvent {
	r.mu.Lock()
	defer r.mu.Unlock()
	return append([]drivercontract.ParsedEvent(nil), r.events...)
}

func TestKimiProcessIdleDeliveryRequiresExactTurnBeginBeforeTerminal(t *testing.T) {
	wireDriver := newKimiTestDriver(16)
	driver := wireDriver
	store := &memoryDeliveryOwnershipStore{}
	recorder := &kimiProcessEventRecorder{}
	process := newKimiProcessDeliveryTestProcess(t, driver, store, true, recorder.record)
	require.True(t, driver.SupportsStdinNotification())
	_, hasStructuredDelivery := any(driver).(drivercontract.RuntimeInputDeliveryDriver)
	require.True(t, hasStructuredDelivery)
	_, hasProtocolDelivery := any(driver).(drivercontract.BusyStdinMessageDriver)
	require.True(t, hasProtocolDelivery, "idle Kimi process delivery must select the prompt RPC path")

	promptStarted := make(chan kimiProcessRPCRequest, 1)
	releasePrompt := make(chan struct{})
	wireDriver.requestFuncForTest = func(ctx context.Context, method string, params any) (json.RawMessage, error) {
		promptStarted <- kimiProcessRPCRequest{method: method, params: params}
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-releasePrompt:
			return json.RawMessage(`{"status":"finished"}`), nil
		}
	}

	deliveryDone := make(chan error, 1)
	go func() {
		deliveryDone <- process.DeliverForHandoff(context.Background(), types.InboxMessage{
			MessageID: "kimi-idle-1",
			Target:    "#agentteam:thread",
			Content:   "handle the idle follow-up",
		})
	}()

	var request kimiProcessRPCRequest
	select {
	case request = <-promptStarted:
	case err := <-deliveryDone:
		t.Fatalf("idle delivery completed without entering the Kimi prompt RPC: %v", err)
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for the Kimi prompt RPC")
	}
	require.Equal(t, "prompt", request.method)
	params, ok := request.params.(map[string]any)
	require.True(t, ok)
	userInput, ok := params["user_input"].(string)
	require.True(t, ok)
	require.Contains(t, userInput, "handle the idle follow-up")
	require.Contains(t, userInput, kimiRuntimeInputMarker("kimi-idle-1"))
	requireKimiOwnershipInflight(t, store, "kimi-idle-1", DeliveryInflightDispatching)

	// Merely printing the marker as assistant content is not correlated
	// model-observation evidence. Only the exact marker echoed by TurnBegin may
	// transition the pending owner.
	wireDriver.handleNotification("event", json.RawMessage(`{
		"type":"ContentPart",
		"payload":{"type":"text","text":"<!-- agentteam-runtime-input:kimi-idle-1 -->"}
	}`))
	routeKimiProcessDriverEvents(t, wireDriver, process, 1)
	requireKimiOwnershipInflight(t, store, "kimi-idle-1", DeliveryInflightDispatching)
	require.Empty(t, kimiProcessEventsOfKind(recorder.snapshot(), consts.ParsedEventKindRuntimeInputAcked))

	wireDriver.handleNotification("event", json.RawMessage(`{
		"type":"TurnBegin",
		"payload":{"user_input":"handle the idle follow-up\n\n<!-- agentteam-runtime-input:kimi-idle-1 -->"}
	}`))
	parsed := routeKimiProcessDriverEvents(t, wireDriver, process, 2)
	require.Equal(t, consts.ParsedEventKindThinking, parsed[0].Kind)
	require.Equal(t, consts.ParsedEventKindRuntimeInputAcked, parsed[1].Kind)
	require.Equal(t, "kimi-idle-1", parsed[1].RuntimeInputID)

	// The exact TurnBegin is the protocol's applied boundary. Delivery returns
	// without waiting for the potentially long-running prompt RPC result.
	select {
	case err := <-deliveryDone:
		require.NoError(t, err)
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for idle Kimi delivery acceptance")
	}
	requireKimiOwnershipAcked(t, store, "kimi-idle-1")
	require.Len(t, kimiProcessEventsOfKind(recorder.snapshot(), consts.ParsedEventKindRuntimeInputAcked), 1)

	wireDriver.handleNotification("event", json.RawMessage(`{"type":"TurnEnd","payload":{}}`))
	select {
	case line := <-wireDriver.lines:
		t.Fatalf("Kimi published a terminal event before the prompt RPC result: %s", line)
	default:
	}
	require.Empty(t, kimiProcessEventsOfKind(recorder.snapshot(), consts.ParsedEventKindTurnEnd))

	close(releasePrompt)
	requireKimiOwnershipAcked(t, store, "kimi-idle-1")
	acked := kimiProcessEventsOfKind(recorder.snapshot(), consts.ParsedEventKindRuntimeInputAcked)
	require.Len(t, acked, 1)
	require.Equal(t, "kimi-idle-1", acked[0].RuntimeInputID)

	terminal := routeKimiProcessDriverEvents(t, wireDriver, process, 1)
	require.Equal(t, consts.ParsedEventKindTurnEnd, terminal[0].Kind)
	require.Equal(t, consts.ParsedEventStopReasonEndTurn, terminal[0].StopReason)
	require.True(t, process.IsIdle())
	require.Len(t, kimiProcessEventsOfKind(recorder.snapshot(), consts.ParsedEventKindTurnEnd), 1)
}

func TestKimiProcessActiveSteerAcceptanceWaitsForExactSteerInput(t *testing.T) {
	driver := newKimiTestDriver(8)
	driver.turnActive = true
	driver.turnTerminal = make(chan struct{})
	store := &memoryDeliveryOwnershipStore{}
	recorder := &kimiProcessEventRecorder{}
	process := newKimiProcessDeliveryTestProcess(t, driver, store, false, recorder.record)

	requested := make(chan kimiProcessRPCRequest, 1)
	driver.requestFuncForTest = func(_ context.Context, method string, params any) (json.RawMessage, error) {
		requested <- kimiProcessRPCRequest{method: method, params: params}
		return json.RawMessage(`{"status":"steered"}`), nil
	}

	deliveryDone := make(chan error, 1)
	go func() {
		deliveryDone <- process.DeliverForHandoff(context.Background(), types.InboxMessage{
			MessageID: "kimi-steer-1",
			Target:    "#agentteam:thread",
			Content:   "apply this while active",
		})
	}()
	var request kimiProcessRPCRequest
	select {
	case request = <-requested:
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for Kimi steer request")
	}
	require.Equal(t, "steer", request.method)
	params, ok := request.params.(map[string]any)
	require.True(t, ok)
	userInput, ok := params["user_input"].(string)
	require.True(t, ok)
	require.Contains(t, userInput, "apply this while active")
	require.Contains(t, userInput, kimiRuntimeInputMarker("kimi-steer-1"))

	// status=steered is transport acceptance, not the applied acknowledgement.
	requireKimiOwnershipInflight(t, store, "kimi-steer-1", DeliveryInflightDispatching)
	require.Empty(t, kimiProcessEventsOfKind(recorder.snapshot(), consts.ParsedEventKindRuntimeInputAcked))
	process.mu.Lock()
	_, pending := process.pendingRuntimeInputs["kimi-steer-1"]
	_, acked := process.ackedSet["kimi-steer-1"]
	process.mu.Unlock()
	require.True(t, pending)
	require.False(t, acked)

	driver.handleNotification("event", json.RawMessage(`{
		"type":"SteerInput",
		"payload":{"user_input":"some other input\n\n<!-- agentteam-runtime-input:kimi-steer-other -->"}
	}`))
	parsed := routeKimiProcessDriverEvents(t, driver, process, 1)
	require.Equal(t, "kimi-steer-other", parsed[0].RuntimeInputID)
	requireKimiOwnershipInflight(t, store, "kimi-steer-1", DeliveryInflightDispatching)
	require.Empty(t, kimiProcessEventsOfKind(recorder.snapshot(), consts.ParsedEventKindRuntimeInputAcked))

	driver.handleNotification("event", json.RawMessage(`{
		"type":"SteerInput",
		"payload":{"user_input":"apply this while active\n\n<!-- agentteam-runtime-input:kimi-steer-1 -->"}
	}`))
	parsed = routeKimiProcessDriverEvents(t, driver, process, 1)
	require.Equal(t, "kimi-steer-1", parsed[0].RuntimeInputID)
	select {
	case err := <-deliveryDone:
		require.NoError(t, err)
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for applied Kimi steer delivery")
	}
	requireKimiOwnershipAcked(t, store, "kimi-steer-1")

	process.mu.Lock()
	_, pending = process.pendingRuntimeInputs["kimi-steer-1"]
	_, acked = process.ackedSet["kimi-steer-1"]
	process.mu.Unlock()
	require.False(t, pending)
	require.True(t, acked)
	ackedEvents := kimiProcessEventsOfKind(recorder.snapshot(), consts.ParsedEventKindRuntimeInputAcked)
	require.Len(t, ackedEvents, 1)
	require.Equal(t, "kimi-steer-1", ackedEvents[0].RuntimeInputID)
}

func newKimiProcessDeliveryTestProcess(
	t *testing.T,
	driver drivercontract.Driver,
	store DeliveryOwnershipStore,
	idle bool,
	onEvent func(drivercontract.ParsedEvent),
) *Process {
	t.Helper()
	exited := make(chan struct{})
	process := &Process{
		AgentID:   "agent-kimi-process-delivery",
		SessionID: kimiTestSessionID,
		processRuntimeBinding: processRuntimeBinding{
			driver:    driver,
			stdinPipe: nopWriteCloser{},
			cancel:    func() {},
			onEvent:   onEvent,
		},
		processLifecycleState: processLifecycleState{
			state:    processStateRunning,
			isIdle:   idle,
			exitCode: -1,
			exited:   exited,
		},
		processDeliveryState: processDeliveryState{
			deliveryOwnershipStore:       store,
			pendingReconcileSuccessGrace: time.Hour,
		},
		processRecoveryState: processRecoveryState{
			workDir:         t.TempDir(),
			disableRecovery: true,
		},
	}
	t.Cleanup(func() {
		stopInboxRetryTestProcess(process)
		close(exited)
	})
	return process
}

func routeKimiProcessDriverEvents(
	t *testing.T,
	driver *KimiWireDriver,
	process *Process,
	count int,
) []drivercontract.ParsedEvent {
	t.Helper()
	events := make([]drivercontract.ParsedEvent, 0, count)
	for len(events) < count {
		select {
		case line := <-driver.lines:
			parsed := driver.ParseLine(line)
			require.NotEmpty(t, parsed)
			for _, event := range parsed {
				process.handleParsedEvent(event)
				events = append(events, event)
			}
		case <-time.After(time.Second):
			t.Fatalf("timed out waiting for Kimi process event %d/%d", len(events), count)
		}
	}
	return events
}

func requireKimiOwnershipAcked(
	t *testing.T,
	store *memoryDeliveryOwnershipStore,
	messageID string,
) {
	t.Helper()
	record, ok := store.record(messageID)
	require.True(t, ok)
	require.True(t, record.IsAcked())
}

func requireKimiOwnershipInflight(t *testing.T, store *memoryDeliveryOwnershipStore, messageID string, reason DeliveryOwnershipInflightReason) {
	t.Helper()
	record, ok := store.record(messageID)
	require.True(t, ok)
	require.True(t, record.IsInflight(reason))
}

func kimiProcessEventsOfKind(events []drivercontract.ParsedEvent, kind string) []drivercontract.ParsedEvent {
	filtered := make([]drivercontract.ParsedEvent, 0, len(events))
	for _, event := range events {
		if event.Kind == kind {
			filtered = append(filtered, event)
		}
	}
	return filtered
}