package agent

import (
	"context"
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"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/internal/core/rewindcaps"
)

// writePiSessionFile writes a Pi-style JSONL session file with the given
// message entries and returns the boundary (pre-message) entry id.
func writePiSessionFile(t *testing.T, dir string, userMessage string) (string, string) {
	t.Helper()
	sessionID := "sess-0001"
	path := filepath.Join(dir, "2026-08-17T00-00-00Z_"+sessionID+".jsonl")
	header, _ := json.Marshal(map[string]any{"type": "session", "version": 3, "id": sessionID})
	info, _ := json.Marshal(map[string]any{"type": "session_info", "id": "info-1", "parentId": nil})
	model, _ := json.Marshal(map[string]any{"type": "model_change", "id": "model-1", "parentId": "info-1"})
	thinking, _ := json.Marshal(map[string]any{"type": "thinking_level_change", "id": "think-1", "parentId": "model-1"})
	userEntry, _ := json.Marshal(map[string]any{
		"type":     "message",
		"id":       "msg-1",
		"parentId": "think-1",
		"message":  map[string]any{"role": "user", "content": []map[string]any{{"type": "text", "text": userMessage}}},
	})
	assistantEntry, _ := json.Marshal(map[string]any{
		"type":     "message",
		"id":       "msg-2",
		"parentId": "msg-1",
		"message":  map[string]any{"role": "assistant", "content": []map[string]any{{"type": "text", "text": "response"}}},
	})
	lines := []string{
		string(header), string(info), string(model), string(thinking),
		string(userEntry), string(assistantEntry),
	}
	require.NoError(t, os.WriteFile(path, []byte(strings.Join(lines, "\n")+"\n"), 0o600))
	// Boundary is the entry before the user message: the thinking_level_change
	// entry (think-1).
	return path, sessionID
}

func TestPiTruncateSessionAtBoundary(t *testing.T) {
	dir := t.TempDir()
	path, _ := writePiSessionFile(t, dir, "first message")

	err := piTruncateSessionAtBoundary(context.Background(), path, "think-1", rewindcaps.RecoveryBoundaryResumeAt)
	require.NoError(t, err)

	data, err := os.ReadFile(path)
	require.NoError(t, err)
	lines := strings.Split(strings.TrimSpace(string(data)), "\n")
	// header + info + model + thinking only (4 lines); user+assistant dropped.
	require.Len(t, lines, 4)
	for _, line := range lines {
		var entry struct {
			Type string `json:"type"`
			ID   string `json:"id"`
		}
		require.NoError(t, json.Unmarshal([]byte(line), &entry))
		require.NotEqual(t, "message", entry.Type, "later context must be dropped")
	}
}

func TestPiTruncateSessionAtBoundaryBeforeFirstInput(t *testing.T) {
	dir := t.TempDir()
	path, _ := writePiSessionFile(t, dir, "first message")

	// before_first_input keeps the header + non-message metadata only.
	err := piTruncateSessionAtBoundary(context.Background(), path, "", rewindcaps.RecoveryBoundaryBeforeFirstInput)
	require.NoError(t, err)

	data, err := os.ReadFile(path)
	require.NoError(t, err)
	lines := strings.Split(strings.TrimSpace(string(data)), "\n")
	for _, line := range lines {
		var entry struct {
			Type string `json:"type"`
		}
		require.NoError(t, json.Unmarshal([]byte(line), &entry))
		require.NotEqual(t, "message", entry.Type)
	}
}

func TestPiTruncateSessionAtBoundaryMissingEntry(t *testing.T) {
	dir := t.TempDir()
	path, _ := writePiSessionFile(t, dir, "first message")
	err := piTruncateSessionAtBoundary(context.Background(), path, "not-present", rewindcaps.RecoveryBoundaryResumeAt)
	require.Error(t, err)
	require.Contains(t, err.Error(), "not found")
}

func TestPiSessionFilePath(t *testing.T) {
	stateDir := t.TempDir()
	sessionDir := filepath.Join(stateDir, ".pi-sessions")
	require.NoError(t, os.MkdirAll(sessionDir, 0o700))
	path, sessionID := writePiSessionFile(t, sessionDir, "hello")
	cfg := drivercontract.LaunchConfig{
		RuntimeStateDir: stateDir,
		WorkDir:         stateDir,
	}
	resolved, err := piSessionFilePath(cfg, sessionID)
	require.NoError(t, err)
	require.Equal(t, path, resolved)
}

func TestPiRPCRecoveryRefRoundTrip(t *testing.T) {
	sessionID := "sess-0001"
	ref, err := rewindcaps.EncodePiSessionTreeRewindRef(sessionID, "entry-1", rewindcaps.RecoveryBoundaryResumeAt)
	require.NoError(t, err)
	decoded, err := rewindcaps.DecodePiSessionTreeRewindRef(ref)
	require.NoError(t, err)
	require.Equal(t, sessionID, decoded.SessionID)
	require.Equal(t, "entry-1", decoded.BoundaryEntryID)
	require.Equal(t, rewindcaps.RecoveryBoundaryResumeAt, decoded.BoundaryKind)
}

func TestPiRPCRecoveryRefValidation(t *testing.T) {
	ref, err := rewindcaps.EncodePiSessionTreeRewindRef("sess-1", "entry-1", rewindcaps.RecoveryBoundaryResumeAt)
	require.NoError(t, err)
	require.NoError(t, rewindcaps.ValidatePreparedRecoveryPoint(rewindcaps.StrategyPiSessionTreeRewind, "sess-1", ref))

	missingBoundary, _ := rewindcaps.EncodePiSessionTreeRewindRef("sess-1", "", rewindcaps.RecoveryBoundaryResumeAt)
	require.Error(t, rewindcaps.ValidatePreparedRecoveryPoint(rewindcaps.StrategyPiSessionTreeRewind, "sess-1", missingBoundary))

	// before_first_input intentionally has no boundary entry.
	fresh, _ := rewindcaps.EncodePiSessionTreeRewindRef("sess-1", "", rewindcaps.RecoveryBoundaryBeforeFirstInput)
	require.NoError(t, rewindcaps.ValidatePreparedRecoveryPoint(rewindcaps.StrategyPiSessionTreeRewind, "sess-1", fresh))

	// Session id mismatch is rejected.
	require.Error(t, rewindcaps.ValidatePreparedRecoveryPoint(rewindcaps.StrategyPiSessionTreeRewind, "other-sess", ref))
}

func TestPiRPCBoundaryKindSelection(t *testing.T) {
	// Fresh session (no durable session file yet) => before_first_input.
	d := &PiRPCDriver{}
	d.mu.Lock()
	d.sessionID = "sess-boundary"
	d.launchCfg = drivercontract.LaunchConfig{
		RuntimeStateDir: t.TempDir(),
		WorkDir:         t.TempDir(),
	}
	d.mu.Unlock()
	strategy, ref, supported, err := d.PrepareRecoveryPointForMessage("msg-1")
	require.NoError(t, err)
	require.True(t, supported)
	require.Equal(t, PiRPCRecoveryStrategy, strategy)
	decoded, err := rewindcaps.DecodePiSessionTreeRewindRef(ref)
	require.NoError(t, err)
	require.Equal(t, rewindcaps.RecoveryBoundaryBeforeFirstInput, decoded.BoundaryKind)
	require.Empty(t, decoded.BoundaryEntryID)
}

func TestPiRPCDriverImplementations(t *testing.T) {
	require.Implements(t, (*drivercontract.Driver)(nil), &PiRPCDriver{})
	require.Implements(t, (*drivercontract.RuntimeInputDeliveryDriver)(nil), &PiRPCDriver{})
	require.Implements(t, (*drivercontract.AckedBusyRuntimeInputDriver)(nil), &PiRPCDriver{})
	require.Implements(t, (*drivercontract.InterruptingRuntimeInputDriver)(nil), &PiRPCDriver{})
	require.Implements(t, (*drivercontract.ActiveRunInterruptDriver)(nil), &PiRPCDriver{})
	require.Implements(t, (*drivercontract.ExplicitTurnLifecycleDriver)(nil), &PiRPCDriver{})
	require.Implements(t, (*drivercontract.RecoveryDriver)(nil), &PiRPCDriver{})
	require.Implements(t, (*drivercontract.MessageBoundRecoveryDriver)(nil), &PiRPCDriver{})
	require.Implements(t, (*drivercontract.DetachedRecoveryDriver)(nil), &PiRPCDriver{})
	require.NotImplements(t, (*drivercontract.IdleCompactionDriver)(nil), &PiRPCDriver{})
	require.Implements(t, (*drivercontract.BusyStdinDeliveryDriver)(nil), &PiRPCDriver{})
	require.Implements(t, (*drivercontract.BusyStdinMessageDriver)(nil), &PiRPCDriver{})
	require.Implements(t, (*drivercontract.CorrelatedRuntimeInputDeliveryDriver)(nil), &PiRPCDriver{})
	_ = runtimeconfig.TypePi
}