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

func TestKimiPrepareRecoveryPointUsesCompletedTurnBoundary(t *testing.T) {
	workDir := t.TempDir()
	shareDir := t.TempDir()
	sessionDir := kimiTestSessionDir(t, shareDir, workDir, kimiTestSessionID)
	require.NoError(t, os.MkdirAll(sessionDir, 0o700))
	writeKimiLines(t, filepath.Join(sessionDir, "wire.jsonl"), []string{
		`{"type":"metadata","protocol_version":"1.10"}`,
		kimiTestWireRecord("TurnBegin", map[string]any{"user_input": "one"}),
		kimiTestWireRecord("SteerInput", map[string]any{"user_input": "one steer"}),
		kimiTestWireRecord("TurnEnd", map[string]any{}),
		kimiTestWireRecord("TurnBegin", map[string]any{"user_input": "two"}),
		kimiTestWireRecord("TurnEnd", map[string]any{}),
		kimiTestWireRecord("TurnBegin", map[string]any{"user_input": "three"}),
		kimiTestWireRecord("TurnEnd", map[string]any{}),
	})
	driver := &KimiWireDriver{
		sessionID:      kimiTestSessionID,
		completedTurns: 3,
		launchCfg: drivercontract.LaunchConfig{
			WorkDir: workDir,
			Env:     map[string]string{"KIMI_SHARE_DIR": shareDir},
		},
	}
	strategy, recoveryRef, supported, err := driver.PrepareRecoveryPoint()
	require.NoError(t, err)
	require.True(t, supported)
	require.Equal(t, rewindcaps.StrategyKimiSessionTurns, strategy)
	require.JSONEq(t, `{"sessionId":"`+kimiTestSessionID+`","turnCount":3,"inputCount":4}`, recoveryRef)
}

func TestKimiPrepareRecoveryPointFailsWithoutDurableWireHistory(t *testing.T) {
	sessionDir := t.TempDir()
	_, err := kimiInputCountAtCompletedTurn(sessionDir, 1)
	require.ErrorIs(t, err, os.ErrNotExist)
}

func TestKimiRestoreRecoveryPointDetachedTruncatesWireAndContext(t *testing.T) {
	for _, target := range []int{0, 2} {
		t.Run(string(rune('0'+target)), func(t *testing.T) {
			workDir := t.TempDir()
			shareDir := t.TempDir()
			sessionDir := kimiTestSessionDir(t, shareDir, workDir, kimiTestSessionID)
			require.NoError(t, os.MkdirAll(sessionDir, 0o700))
			wireLines := []string{`{"type":"metadata","protocol_version":"1.10"}`}
			contextLines := []string{`{"role":"_system_prompt","content":"system"}`}
			for index, text := range []string{"one", "two", "three"} {
				wireLines = append(wireLines,
					kimiTestWireRecord("TurnBegin", map[string]any{"user_input": text}),
					kimiTestWireRecord("ContentPart", map[string]any{"type": "text", "text": text}),
					kimiTestWireRecord("TurnEnd", map[string]any{}),
				)
				contextLines = append(contextLines,
					mustKimiJSON(map[string]any{"role": "user", "content": text}),
					mustKimiJSON(map[string]any{"role": "assistant", "content": "reply"}),
				)
				if index == 0 {
					contextLines = append(contextLines, `{"role":"user","content":"<system>CHECKPOINT 0</system>"}`)
				}
			}
			writeKimiLines(t, filepath.Join(sessionDir, "wire.jsonl"), wireLines)
			writeKimiLines(t, filepath.Join(sessionDir, "context.jsonl"), contextLines)
			require.NoError(t, os.WriteFile(filepath.Join(sessionDir, "state.json"), []byte(`{
				"title":"retained",
				"approval":{"yolo":true},
				"plan_mode":true,
				"plan_session_id":"plan-1",
				"plan_slug":"plan-file",
				"todos":[{"title":"later","status":"in_progress"}]
			}`), 0o600))
			recoveryRef, err := rewindcaps.EncodeKimiSessionTurnsRecoveryRef(kimiTestSessionID, target)
			require.NoError(t, err)
			err = (&KimiWireDriver{}).RestoreRecoveryPointDetached(
				context.Background(),
				drivercontract.LaunchConfig{
					WorkDir: workDir,
					Env:     map[string]string{"KIMI_SHARE_DIR": shareDir},
				},
				rewindcaps.StrategyKimiSessionTurns,
				recoveryRef,
			)
			require.NoError(t, err)
			wireData, err := os.ReadFile(filepath.Join(sessionDir, "wire.jsonl"))
			require.NoError(t, err)
			contextData, err := os.ReadFile(filepath.Join(sessionDir, "context.jsonl"))
			require.NoError(t, err)
			require.Equal(t, target, strings.Count(string(wireData), `"type":"TurnEnd"`))
			require.Equal(t, target, countKimiContextUsers(t, string(contextData)))
			require.Contains(t, string(contextData), `"_system_prompt"`)
			stateData, err := os.ReadFile(filepath.Join(sessionDir, "state.json"))
			require.NoError(t, err)
			require.JSONEq(t, `{"title":"retained","approval":{"yolo":true},"plan_mode":false,"todos":[]}`, string(stateData))
		})
	}
}

func TestKimiRestoreRecoveryPointKeepsSteerInputsInsideRetainedTurn(t *testing.T) {
	workDir := t.TempDir()
	shareDir := t.TempDir()
	sessionDir := kimiTestSessionDir(t, shareDir, workDir, kimiTestSessionID)
	require.NoError(t, os.MkdirAll(sessionDir, 0o700))
	wireLines := []string{
		`{"type":"metadata","protocol_version":"1.10"}`,
		kimiTestWireRecord("TurnBegin", map[string]any{"user_input": "one"}),
		kimiTestWireRecord("SteerInput", map[string]any{"user_input": "steer one"}),
		kimiTestWireRecord("TurnEnd", map[string]any{}),
		kimiTestWireRecord("TurnBegin", map[string]any{"user_input": "two"}),
		kimiTestWireRecord("TurnEnd", map[string]any{}),
	}
	contextLines := []string{
		`{"role":"_system_prompt","content":"system"}`,
		`{"role":"user","content":"one"}`,
		`{"role":"assistant","content":"partial"}`,
		`{"role":"user","content":"steer one"}`,
		`{"role":"assistant","content":"done"}`,
		`{"role":"user","content":"two"}`,
		`{"role":"assistant","content":"done"}`,
	}
	writeKimiLines(t, filepath.Join(sessionDir, "wire.jsonl"), wireLines)
	writeKimiLines(t, filepath.Join(sessionDir, "context.jsonl"), contextLines)
	require.NoError(t, os.WriteFile(filepath.Join(sessionDir, "state.json"), []byte(`{"plan_mode":false}`), 0o600))
	recoveryRef, err := rewindcaps.EncodeKimiSessionTurnsRecoveryRef(kimiTestSessionID, 1, 2)
	require.NoError(t, err)
	require.NoError(t, (&KimiWireDriver{}).RestoreRecoveryPointDetached(
		context.Background(),
		drivercontract.LaunchConfig{WorkDir: workDir, Env: map[string]string{"KIMI_SHARE_DIR": shareDir}},
		rewindcaps.StrategyKimiSessionTurns,
		recoveryRef,
	))
	contextData, err := os.ReadFile(filepath.Join(sessionDir, "context.jsonl"))
	require.NoError(t, err)
	require.Contains(t, string(contextData), `"content":"steer one"`)
	require.NotContains(t, string(contextData), `"content":"two"`)
}

func TestKimiRestoreRecoveryPointDetachedRejectsMismatchedInputCountWithoutMutation(t *testing.T) {
	workDir := t.TempDir()
	shareDir := t.TempDir()
	sessionDir := kimiTestSessionDir(t, shareDir, workDir, kimiTestSessionID)
	require.NoError(t, os.MkdirAll(sessionDir, 0o700))
	wirePath := filepath.Join(sessionDir, "wire.jsonl")
	contextPath := filepath.Join(sessionDir, "context.jsonl")
	statePath := filepath.Join(sessionDir, "state.json")
	writeKimiLines(t, wirePath, []string{
		`{"type":"metadata","protocol_version":"1.10"}`,
		kimiTestWireRecord("TurnBegin", map[string]any{"user_input": "one"}),
		kimiTestWireRecord("SteerInput", map[string]any{"user_input": "steer one"}),
		kimiTestWireRecord("TurnEnd", map[string]any{}),
		kimiTestWireRecord("TurnBegin", map[string]any{"user_input": "two"}),
		kimiTestWireRecord("TurnEnd", map[string]any{}),
	})
	writeKimiLines(t, contextPath, []string{
		`{"role":"user","content":"one"}`,
		`{"role":"assistant","content":"partial"}`,
		`{"role":"user","content":"steer one"}`,
		`{"role":"assistant","content":"done"}`,
		`{"role":"user","content":"two"}`,
	})
	require.NoError(t, os.WriteFile(statePath, []byte(`{"plan_mode":true}`), 0o600))
	beforeWire, err := os.ReadFile(wirePath)
	require.NoError(t, err)
	beforeContext, err := os.ReadFile(contextPath)
	require.NoError(t, err)
	beforeState, err := os.ReadFile(statePath)
	require.NoError(t, err)
	recoveryRef, err := rewindcaps.EncodeKimiSessionTurnsRecoveryRef(kimiTestSessionID, 1, 1)
	require.NoError(t, err)

	err = (&KimiWireDriver{}).RestoreRecoveryPointDetached(
		context.Background(),
		drivercontract.LaunchConfig{WorkDir: workDir, Env: map[string]string{"KIMI_SHARE_DIR": shareDir}},
		rewindcaps.StrategyKimiSessionTurns,
		recoveryRef,
	)
	require.ErrorContains(t, err, "does not match durable wire history")
	afterWire, err := os.ReadFile(wirePath)
	require.NoError(t, err)
	afterContext, err := os.ReadFile(contextPath)
	require.NoError(t, err)
	afterState, err := os.ReadFile(statePath)
	require.NoError(t, err)
	require.Equal(t, beforeWire, afterWire)
	require.Equal(t, beforeContext, afterContext)
	require.Equal(t, beforeState, afterState)
}

func TestKimiRestoreRecoveryPointDetachedFailsClosed(t *testing.T) {
	tests := []struct {
		name      string
		wireLines []string
		context   []string
		target    int
	}{
		{
			name:      "target ahead",
			wireLines: []string{`{"type":"metadata"}`, kimiTestWireRecord("TurnBegin", map[string]any{"user_input": "one"}), kimiTestWireRecord("TurnEnd", map[string]any{})},
			context:   []string{`{"role":"user","content":"one"}`},
			target:    2,
		},
		{
			name:      "incomplete wire turn",
			wireLines: []string{`{"type":"metadata"}`, kimiTestWireRecord("TurnBegin", map[string]any{"user_input": "one"})},
			context:   []string{`{"role":"user","content":"one"}`},
			target:    1,
		},
		{
			name:      "malformed context",
			wireLines: []string{`{"type":"metadata"}`},
			context:   []string{`{bad`},
			target:    0,
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			workDir := t.TempDir()
			shareDir := t.TempDir()
			sessionDir := kimiTestSessionDir(t, shareDir, workDir, kimiTestSessionID)
			require.NoError(t, os.MkdirAll(sessionDir, 0o700))
			wirePath := filepath.Join(sessionDir, "wire.jsonl")
			contextPath := filepath.Join(sessionDir, "context.jsonl")
			writeKimiLines(t, wirePath, test.wireLines)
			writeKimiLines(t, contextPath, test.context)
			beforeWire, _ := os.ReadFile(wirePath)
			beforeContext, _ := os.ReadFile(contextPath)
			recoveryRef, err := rewindcaps.EncodeKimiSessionTurnsRecoveryRef(kimiTestSessionID, test.target)
			require.NoError(t, err)
			err = (&KimiWireDriver{}).RestoreRecoveryPointDetached(
				context.Background(),
				drivercontract.LaunchConfig{WorkDir: workDir, Env: map[string]string{"KIMI_SHARE_DIR": shareDir}},
				rewindcaps.StrategyKimiSessionTurns,
				recoveryRef,
			)
			require.Error(t, err)
			afterWire, _ := os.ReadFile(wirePath)
			afterContext, _ := os.ReadFile(contextPath)
			require.Equal(t, beforeWire, afterWire)
			require.Equal(t, beforeContext, afterContext)
		})
	}
}

func TestKimiRestoreRecoveryPointDetachedCreatesSafeDefaultState(t *testing.T) {
	workDir := t.TempDir()
	shareDir := t.TempDir()
	sessionDir := kimiTestSessionDir(t, shareDir, workDir, kimiTestSessionID)
	require.NoError(t, os.MkdirAll(sessionDir, 0o700))
	wirePath := filepath.Join(sessionDir, "wire.jsonl")
	contextPath := filepath.Join(sessionDir, "context.jsonl")
	writeKimiLines(t, wirePath, []string{
		`{"type":"metadata","protocol_version":"1.10"}`,
		kimiTestWireRecord("TurnBegin", map[string]any{"user_input": "one"}),
		kimiTestWireRecord("TurnEnd", map[string]any{}),
	})
	writeKimiLines(t, contextPath, []string{
		`{"role":"user","content":"one"}`,
		`{"role":"assistant","content":"reply"}`,
	})
	recoveryRef, err := rewindcaps.EncodeKimiSessionTurnsRecoveryRef(kimiTestSessionID, 0, 0)
	require.NoError(t, err)

	err = (&KimiWireDriver{}).RestoreRecoveryPointDetached(
		context.Background(),
		drivercontract.LaunchConfig{WorkDir: workDir, Env: map[string]string{"KIMI_SHARE_DIR": shareDir}},
		rewindcaps.StrategyKimiSessionTurns,
		recoveryRef,
	)
	require.NoError(t, err)
	stateData, readErr := os.ReadFile(filepath.Join(sessionDir, "state.json"))
	require.NoError(t, readErr)
	require.JSONEq(t, `{"plan_mode":false,"todos":[]}`, string(stateData))
}

func TestReplaceKimiRecoveryFilesRollsBackWhenInstallFails(t *testing.T) {
	sessionDir := t.TempDir()
	wirePath := filepath.Join(sessionDir, "wire.jsonl")
	contextPath := filepath.Join(sessionDir, "context.jsonl")
	statePath := filepath.Join(sessionDir, "state.json")
	writeKimiLines(t, wirePath, []string{`{"old":"wire"}`})
	writeKimiLines(t, contextPath, []string{`{"old":"context"}`})
	writeKimiLines(t, statePath, []string{`{"old":"state"}`})

	for _, failureCall := range []int{2, 3} {
		t.Run(string(rune('0'+failureCall)), func(t *testing.T) {
			renameCalls := 0
			rename := func(oldPath, newPath string) error {
				renameCalls++
				if renameCalls == failureCall {
					return os.ErrPermission
				}
				return os.Rename(oldPath, newPath)
			}
			err := replaceKimiRecoveryFilesWithRename(
				[]kimiRecoveryReplacement{
					{path: wirePath, content: kimiJSONLBytes([]string{`{"new":"wire"}`})},
					{path: contextPath, content: kimiJSONLBytes([]string{`{"new":"context"}`})},
					{path: statePath, content: kimiJSONLBytes([]string{`{"new":"state"}`})},
				},
				rename,
			)
			require.ErrorIs(t, err, os.ErrPermission)

			wireData, readErr := os.ReadFile(wirePath)
			require.NoError(t, readErr)
			contextData, readErr := os.ReadFile(contextPath)
			require.NoError(t, readErr)
			stateData, readErr := os.ReadFile(statePath)
			require.NoError(t, readErr)
			require.JSONEq(t, `{"old":"wire"}`, strings.TrimSpace(string(wireData)))
			require.JSONEq(t, `{"old":"context"}`, strings.TrimSpace(string(contextData)))
			require.JSONEq(t, `{"old":"state"}`, strings.TrimSpace(string(stateData)))
			require.NoFileExists(t, filepath.Join(sessionDir, ".agentteam-recovery.json"))
			require.NoFileExists(t, wirePath+".agentteam-backup")
			require.NoFileExists(t, contextPath+".agentteam-backup")
			require.NoFileExists(t, statePath+".agentteam-backup")
		})
		if failureCall == 2 {
			writeKimiLines(t, wirePath, []string{`{"old":"wire"}`})
			writeKimiLines(t, contextPath, []string{`{"old":"context"}`})
			writeKimiLines(t, statePath, []string{`{"old":"state"}`})
		}
	}
}

func TestRecoverKimiRecoveryTransactionRollsBackInterruptedInstall(t *testing.T) {
	sessionDir := t.TempDir()
	wirePath := filepath.Join(sessionDir, "wire.jsonl")
	contextPath := filepath.Join(sessionDir, "context.jsonl")
	statePath := filepath.Join(sessionDir, "state.json")
	writeKimiLines(t, wirePath, []string{`{"new":"wire"}`})
	writeKimiLines(t, contextPath, []string{`{"old":"context"}`})
	writeKimiLines(t, statePath, []string{`{"new":"state"}`})
	writeKimiLines(t, wirePath+".agentteam-backup", []string{`{"old":"wire"}`})
	writeKimiLines(t, contextPath+".agentteam-backup", []string{`{"old":"context"}`})
	writeKimiLines(t, statePath+".agentteam-backup", []string{`{"old":"state"}`})
	require.NoError(t, os.WriteFile(filepath.Join(sessionDir, ".agentteam-recovery.json"), []byte(`{
		"version":1,
		"files":[
			{"path":"wire.jsonl","backup":"wire.jsonl.agentteam-backup"},
			{"path":"context.jsonl","backup":"context.jsonl.agentteam-backup"},
			{"path":"state.json","backup":"state.json.agentteam-backup"}
		]
	}`), 0o600))
	require.NoError(t, recoverKimiRecoveryTransaction(sessionDir))
	wireData, err := os.ReadFile(wirePath)
	require.NoError(t, err)
	contextData, err := os.ReadFile(contextPath)
	require.NoError(t, err)
	stateData, err := os.ReadFile(statePath)
	require.NoError(t, err)
	require.JSONEq(t, `{"old":"wire"}`, strings.TrimSpace(string(wireData)))
	require.JSONEq(t, `{"old":"context"}`, strings.TrimSpace(string(contextData)))
	require.JSONEq(t, `{"old":"state"}`, strings.TrimSpace(string(stateData)))
	require.NoFileExists(t, filepath.Join(sessionDir, ".agentteam-recovery.json"))
	require.NoFileExists(t, wirePath+".agentteam-backup")
	require.NoFileExists(t, contextPath+".agentteam-backup")
	require.NoFileExists(t, statePath+".agentteam-backup")
}

func TestRecoverKimiRecoveryTransactionCanResumeInterruptedRollback(t *testing.T) {
	sessionDir := t.TempDir()
	paths := []string{
		filepath.Join(sessionDir, "wire.jsonl"),
		filepath.Join(sessionDir, "context.jsonl"),
		filepath.Join(sessionDir, "state.json"),
	}
	for _, path := range paths {
		writeKimiLines(t, path, []string{`{"value":"new"}`})
		writeKimiLines(t, path+".agentteam-backup", []string{`{"value":"old"}`})
	}
	require.NoError(t, os.WriteFile(filepath.Join(sessionDir, ".agentteam-recovery.json"), []byte(`{
		"version":1,
		"files":[
			{"path":"wire.jsonl","backup":"wire.jsonl.agentteam-backup"},
			{"path":"context.jsonl","backup":"context.jsonl.agentteam-backup"},
			{"path":"state.json","backup":"state.json.agentteam-backup"}
		]
	}`), 0o600))
	renameCalls := 0
	err := recoverKimiRecoveryTransactionWithRename(sessionDir, func(oldPath, newPath string) error {
		renameCalls++
		if renameCalls == 2 {
			return os.ErrPermission
		}
		return os.Rename(oldPath, newPath)
	})
	require.ErrorIs(t, err, os.ErrPermission)
	require.FileExists(t, filepath.Join(sessionDir, ".agentteam-recovery.json"))
	for _, path := range paths {
		require.FileExists(t, path+".agentteam-backup")
	}

	require.NoError(t, recoverKimiRecoveryTransaction(sessionDir))
	for _, path := range paths {
		data, readErr := os.ReadFile(path)
		require.NoError(t, readErr)
		require.JSONEq(t, `{"value":"old"}`, strings.TrimSpace(string(data)))
		require.NoFileExists(t, path+".agentteam-backup")
	}
	require.NoFileExists(t, filepath.Join(sessionDir, ".agentteam-recovery.json"))
}

func TestKimiSessionDirRejectsPathTraversal(t *testing.T) {
	workDir := t.TempDir()
	_, err := kimiSessionDir(drivercontract.LaunchConfig{
		WorkDir: workDir,
		Env:     map[string]string{"KIMI_SHARE_DIR": t.TempDir()},
	}, "../escape", false)
	require.Error(t, err)
}

func TestKimiSessionDirRejectsAncestorSymlink(t *testing.T) {
	workDir := t.TempDir()
	shareDir := t.TempDir()
	sessionDir := kimiTestSessionDir(t, shareDir, workDir, kimiTestSessionID)
	outside := t.TempDir()
	require.NoError(t, os.MkdirAll(filepath.Join(outside, kimiTestSessionID), 0o700))
	require.NoError(t, os.MkdirAll(filepath.Dir(filepath.Dir(sessionDir)), 0o700))
	require.NoError(t, os.Symlink(outside, filepath.Dir(sessionDir)))

	_, err := kimiSessionDir(drivercontract.LaunchConfig{
		WorkDir: workDir,
		Env:     map[string]string{"KIMI_SHARE_DIR": shareDir},
	}, kimiTestSessionID, true)
	require.Error(t, err)
}

func writeKimiLines(t *testing.T, path string, lines []string) {
	t.Helper()
	content := ""
	if len(lines) > 0 {
		content = strings.Join(lines, "\n") + "\n"
	}
	require.NoError(t, os.WriteFile(path, []byte(content), 0o600))
}

func mustKimiJSON(value any) string {
	data, _ := json.Marshal(value)
	return string(data)
}

func countKimiContextUsers(t *testing.T, data string) int {
	t.Helper()
	count := 0
	for _, line := range strings.Split(strings.TrimSpace(data), "\n") {
		var record map[string]any
		require.NoError(t, json.Unmarshal([]byte(line), &record))
		if record["role"] == "user" && !kimiCheckpointContextRecord(record) {
			count++
		}
	}
	return count
}