package agent

import (
	"bufio"
	"context"
	"crypto/md5"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	"github.com/google/uuid"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/runtimeconfig"
	"github.com/tidbcloud/AgentTeam/internal/core/rewindcaps"
)

var kimiCheckpointPattern = regexp.MustCompile(`^<system>CHECKPOINT [0-9]+</system>$`)
var kimiRuntimeInputMarkerPattern = regexp.MustCompile(`(?:^|\n\n)<!-- agentteam-runtime-input:([^<>\s]+) -->$`)
var errKimiIncompleteTurn = errors.New("kimi: incomplete active turn")

func kimiCompletedTurnCount(cfg drivercontract.LaunchConfig, sessionID string) (int, error) {
	sessionDir, err := kimiSessionDir(cfg, sessionID, false)
	if err != nil {
		return 0, err
	}
	lines, err := readKimiJSONL(filepath.Join(sessionDir, "wire.jsonl"))
	if errors.Is(err, os.ErrNotExist) {
		return 0, nil
	}
	if err != nil {
		if _, statErr := os.Stat(filepath.Join(sessionDir, "wire.jsonl")); errors.Is(statErr, os.ErrNotExist) {
			return 0, nil
		}
		return 0, err
	}
	_, completed, _, err := kimiWirePrefix(lines, -1)
	if err != nil {
		return 0, err
	}
	return completed, nil
}

func (d *KimiWireDriver) PrepareRecoveryPoint() (string, string, bool, error) {
	d.mu.Lock()
	sessionID := d.sessionID
	turnCount := d.completedTurns
	cfg := cloneKimiLaunchConfig(d.launchCfg)
	d.mu.Unlock()
	if sessionID == "" {
		return "", "", false, nil
	}
	sessionDir, err := kimiSessionDir(cfg, sessionID, false)
	if err != nil {
		return "", "", false, err
	}
	inputCount, countErr := kimiInputCountAtCompletedTurn(sessionDir, turnCount)
	if countErr != nil {
		return "", "", false, countErr
	}
	ref, err := rewindcaps.EncodeKimiSessionTurnsRecoveryRef(sessionID, turnCount, inputCount)
	if err != nil {
		return "", "", false, err
	}
	return rewindcaps.StrategyKimiSessionTurns, ref, true, nil
}

func kimiInputCountAtCompletedTurn(sessionDir string, turnCount int) (int, error) {
	if turnCount < 0 {
		return 0, fmt.Errorf("kimi: invalid completed turn count %d", turnCount)
	}
	wireLines, err := readKimiJSONL(filepath.Join(sessionDir, "wire.jsonl"))
	if errors.Is(err, os.ErrNotExist) {
		if turnCount == 0 {
			return 0, nil
		}
		return 0, fmt.Errorf("kimi: durable wire history is unavailable for completed turn %d: %w", turnCount, os.ErrNotExist)
	}
	if err != nil {
		return 0, err
	}
	_, completed, inputCount, err := kimiWirePrefix(wireLines, turnCount)
	if err != nil {
		return 0, err
	}
	if completed < turnCount {
		return 0, fmt.Errorf("kimi: completed turn count %d is ahead of durable history %d", turnCount, completed)
	}
	return inputCount, nil
}

func (d *KimiWireDriver) RestoreRecoveryPoint(context.Context, string, string) error {
	return fmt.Errorf("kimi: live recovery restore is not supported")
}

func (d *KimiWireDriver) RestoreRecoveryPointDetached(
	ctx context.Context,
	cfg drivercontract.LaunchConfig,
	strategy string,
	recoveryRef string,
) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	if strings.TrimSpace(strategy) != rewindcaps.StrategyKimiSessionTurns {
		return fmt.Errorf("kimi: unsupported recovery strategy %q", strategy)
	}
	ref, err := rewindcaps.DecodeKimiSessionTurnsRecoveryRef(recoveryRef)
	if err != nil {
		return fmt.Errorf("kimi: decode recovery ref: %w", err)
	}
	if ref.SessionID == "" || ref.TurnCount < 0 || ref.InputCount < 0 {
		return fmt.Errorf("kimi: invalid recovery ref")
	}
	cfg, err = kimiSessionLookupConfig(cfg)
	if err != nil {
		return err
	}
	sessionDir, err := kimiSessionDir(cfg, ref.SessionID, true)
	if err != nil {
		return err
	}
	if err := recoverKimiRecoveryTransaction(sessionDir); err != nil {
		return fmt.Errorf("kimi: recover interrupted rewind: %w", err)
	}
	wirePath := filepath.Join(sessionDir, "wire.jsonl")
	contextPath := filepath.Join(sessionDir, "context.jsonl")
	statePath := filepath.Join(sessionDir, "state.json")
	wire, err := readKimiJSONL(wirePath)
	if err != nil {
		return err
	}
	contextLines, err := readKimiJSONL(contextPath)
	if err != nil {
		return err
	}
	truncatedWire, completed, wireInputCount, err := kimiWirePrefix(wire, ref.TurnCount)
	if err != nil {
		return err
	}
	if completed < ref.TurnCount {
		return fmt.Errorf("kimi: recovery target %d is ahead of completed turn count %d", ref.TurnCount, completed)
	}
	if wireInputCount != ref.InputCount {
		return fmt.Errorf(
			"kimi: recovery input count %d does not match durable wire history %d at turn %d",
			ref.InputCount,
			wireInputCount,
			ref.TurnCount,
		)
	}
	truncatedContext, err := truncateKimiContext(contextLines, ref.InputCount)
	if err != nil {
		return err
	}
	if contextRealUserCount(contextLines) < ref.InputCount {
		return fmt.Errorf("kimi: context history has fewer user inputs than recovery target %d", ref.InputCount)
	}
	truncatedState, err := resetKimiRecoveryState(statePath)
	if err != nil {
		return err
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	return replaceKimiRecoveryFiles(
		[]kimiRecoveryReplacement{
			{path: wirePath, content: kimiJSONLBytes(truncatedWire)},
			{path: contextPath, content: kimiJSONLBytes(truncatedContext)},
			{path: statePath, content: truncatedState},
		},
	)
}

func resetKimiRecoveryState(path string) ([]byte, error) {
	data, err := os.ReadFile(path)
	if errors.Is(err, os.ErrNotExist) {
		data = []byte(`{}`)
	} else if err != nil {
		return nil, fmt.Errorf("kimi: read state.json: %w", err)
	}
	var state map[string]any
	if err := json.Unmarshal(data, &state); err != nil {
		return nil, fmt.Errorf("kimi: malformed state.json: %w", err)
	}
	reset := map[string]any{
		"plan_mode": false,
		"todos":     []any{},
	}
	for _, key := range []string{"title", "approval"} {
		if value, ok := state[key]; ok {
			reset[key] = value
		}
	}
	payload, err := json.Marshal(reset)
	if err != nil {
		return nil, fmt.Errorf("kimi: encode state.json: %w", err)
	}
	return append(payload, '\n'), nil
}

func (d *KimiWireDriver) ReconcileRuntimeInput(
	ctx context.Context,
	inputID string,
	sessionID string,
) (RuntimeInputReconciliation, error) {
	if err := ctx.Err(); err != nil {
		return RuntimeInputReconciliationUnknown, err
	}
	d.mu.Lock()
	cfg := cloneKimiLaunchConfig(d.launchCfg)
	d.mu.Unlock()
	sessionDir, err := kimiSessionDir(cfg, sessionID, true)
	if err != nil {
		return RuntimeInputReconciliationUnknown, err
	}
	lines, err := readKimiJSONL(filepath.Join(sessionDir, "wire.jsonl"))
	if err != nil {
		return RuntimeInputReconciliationUnknown, err
	}
	inputID = strings.TrimSpace(inputID)
	if inputID == "" {
		return RuntimeInputReconciliationUnknown, fmt.Errorf("kimi: runtime input id is required")
	}
	active := false
	sawTurn := false
	for _, line := range lines {
		messageType, payload := kimiWireRecordFields(line)
		switch messageType {
		case "TurnBegin":
			if active {
				return RuntimeInputReconciliationUnknown, fmt.Errorf("kimi: malformed nested turn history")
			}
			active = true
			sawTurn = true
			if kimiRuntimeInputIDFromPayload(payload) == inputID {
				return RuntimeInputReconciliationAcked, nil
			}
		case "SteerInput":
			if !active {
				return RuntimeInputReconciliationUnknown, fmt.Errorf("kimi: malformed SteerInput without active turn")
			}
			if kimiRuntimeInputIDFromPayload(payload) == inputID {
				return RuntimeInputReconciliationAcked, nil
			}
		case "TurnEnd":
			if !active {
				return RuntimeInputReconciliationUnknown, fmt.Errorf("kimi: malformed terminal history")
			}
			active = false
		}
	}
	if active {
		return RuntimeInputReconciliationPending, nil
	}
	if !sawTurn {
		return RuntimeInputReconciliationUnknown, nil
	}
	return RuntimeInputReconciliationAbsent, nil
}

func kimiRuntimeInputMarker(inputID string) string {
	return "<!-- agentteam-runtime-input:" + strings.TrimSpace(inputID) + " -->"
}

func kimiSessionDir(cfg drivercontract.LaunchConfig, sessionID string, requireExists bool) (string, error) {
	if _, err := uuid.Parse(strings.TrimSpace(sessionID)); err != nil {
		return "", fmt.Errorf("kimi: invalid session id %q", sessionID)
	}
	workDir, err := filepath.Abs(strings.TrimSpace(cfg.WorkDir))
	if err != nil || strings.TrimSpace(cfg.WorkDir) == "" {
		return "", fmt.Errorf("kimi: valid workdir is required")
	}
	workDir, err = filepath.EvalSymlinks(filepath.Clean(workDir))
	if err != nil {
		return "", fmt.Errorf("kimi: resolve workdir: %w", err)
	}
	env := runtimeconfig.CloneStringMap(cfg.Env)
	shareDir, err := resolveKimiShareDir(env, workDir)
	if err != nil {
		return "", fmt.Errorf("kimi: resolve share directory: %w", err)
	}
	sum := md5.Sum([]byte(workDir))
	sessionsRoot := filepath.Join(shareDir, "sessions", hex.EncodeToString(sum[:]))
	sessionDir := filepath.Join(sessionsRoot, sessionID)
	relative, err := filepath.Rel(sessionsRoot, sessionDir)
	if err != nil || relative == ".." || strings.HasPrefix(relative, ".."+string(filepath.Separator)) {
		return "", fmt.Errorf("kimi: unsafe session path")
	}
	if requireExists {
		if err := validateKimiDirectoryChain(shareDir, sessionDir); err != nil {
			return "", fmt.Errorf("kimi: session directory unavailable: %w", err)
		}
	}
	return sessionDir, nil
}

func readKimiJSONL(path string) ([]string, error) {
	info, err := os.Lstat(path)
	if err != nil {
		return nil, fmt.Errorf("kimi: inspect %s: %w", filepath.Base(path), err)
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
		return nil, fmt.Errorf("kimi: %s is not a safe regular file", filepath.Base(path))
	}
	file, err := os.Open(path)
	if err != nil {
		return nil, fmt.Errorf("kimi: open %s: %w", filepath.Base(path), err)
	}
	defer file.Close()
	var lines []string
	scanner := bufio.NewScanner(file)
	scanner.Buffer(make([]byte, 64*1024), maxProcessLineBytes)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" {
			continue
		}
		var decoded any
		if err := json.Unmarshal([]byte(line), &decoded); err != nil {
			return nil, fmt.Errorf("kimi: malformed %s: %w", filepath.Base(path), err)
		}
		lines = append(lines, line)
	}
	if err := scanner.Err(); err != nil {
		return nil, fmt.Errorf("kimi: read %s: %w", filepath.Base(path), err)
	}
	return lines, nil
}

func kimiWirePrefix(lines []string, turnCount int) ([]string, int, int, error) {
	var output []string
	active := false
	completed := 0
	userInputs := 0
	for _, line := range lines {
		var record map[string]any
		_ = json.Unmarshal([]byte(line), &record)
		if record["type"] == "metadata" {
			output = append(output, line)
			continue
		}
		messageType, _ := kimiWireRecordFields(line)
		switch messageType {
		case "TurnBegin":
			if active {
				return nil, completed, userInputs, fmt.Errorf("kimi: malformed nested TurnBegin")
			}
			if turnCount >= 0 && completed >= turnCount {
				return output, completed, userInputs, nil
			}
			active = true
			userInputs++
		case "SteerInput":
			if !active {
				return nil, completed, userInputs, fmt.Errorf("kimi: SteerInput without active turn")
			}
			userInputs++
		case "TurnEnd":
			if !active {
				return nil, completed, userInputs, fmt.Errorf("kimi: TurnEnd without TurnBegin")
			}
			completed++
			active = false
		}
		if turnCount < 0 || completed < turnCount || active || messageType == "TurnEnd" {
			output = append(output, line)
		}
		if turnCount >= 0 && completed == turnCount && messageType == "TurnEnd" {
			return output, completed, userInputs, nil
		}
	}
	if active {
		return nil, completed, userInputs, errKimiIncompleteTurn
	}
	return output, completed, userInputs, nil
}

func kimiWireRecordFields(line string) (string, json.RawMessage) {
	var record struct {
		Message struct {
			Type    string          `json:"type"`
			Payload json.RawMessage `json:"payload"`
		} `json:"message"`
	}
	if json.Unmarshal([]byte(line), &record) != nil {
		return "", nil
	}
	return record.Message.Type, record.Message.Payload
}

func truncateKimiContext(lines []string, userInputCount int) ([]string, error) {
	var output []string
	users := 0
	for _, line := range lines {
		var record map[string]any
		if err := json.Unmarshal([]byte(line), &record); err != nil {
			return nil, err
		}
		if record["role"] == "user" && !kimiCheckpointContextRecord(record) {
			if users >= userInputCount {
				break
			}
			users++
		}
		if users <= userInputCount {
			output = append(output, line)
		}
	}
	return output, nil
}

func contextRealUserCount(lines []string) int {
	count := 0
	for _, line := range lines {
		var record map[string]any
		if json.Unmarshal([]byte(line), &record) == nil &&
			record["role"] == "user" &&
			!kimiCheckpointContextRecord(record) {
			count++
		}
	}
	return count
}

func kimiCheckpointContextRecord(record map[string]any) bool {
	if record["role"] != "user" {
		return false
	}
	content := record["content"]
	if text, ok := content.(string); ok {
		return kimiCheckpointPattern.MatchString(strings.TrimSpace(text))
	}
	parts, ok := content.([]any)
	if !ok || len(parts) != 1 {
		return false
	}
	part, ok := parts[0].(map[string]any)
	if !ok {
		return false
	}
	text, _ := part["text"].(string)
	return kimiCheckpointPattern.MatchString(strings.TrimSpace(text))
}

type kimiRecoveryReplacement struct {
	path              string
	content           []byte
	originallyMissing bool
}

func replaceKimiRecoveryFiles(files []kimiRecoveryReplacement) error {
	return replaceKimiRecoveryFilesWithRename(files, replaceKimiRecoveryFile)
}

func replaceKimiRecoveryFilesWithRename(
	files []kimiRecoveryReplacement,
	rename func(string, string) error,
) error {
	if len(files) == 0 {
		return fmt.Errorf("kimi: recovery transaction has no files")
	}
	sessionDir := filepath.Dir(files[0].path)
	for _, file := range files {
		if filepath.Dir(file.path) != sessionDir || filepath.Base(file.path) != filepath.Clean(filepath.Base(file.path)) {
			return fmt.Errorf("kimi: recovery files must share one session directory")
		}
	}
	if err := recoverKimiRecoveryTransaction(sessionDir); err != nil {
		return err
	}
	type stagedFile struct {
		path   string
		tmp    string
		backup string
	}
	staged := make([]stagedFile, 0, len(files))
	contents := make([][]byte, 0, len(files))
	for index, file := range files {
		info, statErr := os.Lstat(file.path)
		switch {
		case errors.Is(statErr, os.ErrNotExist):
			files[index].originallyMissing = true
		case statErr != nil:
			return fmt.Errorf("kimi: inspect recovery file %s: %w", filepath.Base(file.path), statErr)
		case info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular():
			return fmt.Errorf("kimi: recovery file %s is not a safe regular file", filepath.Base(file.path))
		}
		staged = append(staged, stagedFile{path: file.path, backup: file.path + ".agentteam-backup"})
		contents = append(contents, file.content)
	}
	for _, file := range staged {
		if err := os.Remove(file.backup); err != nil && !errors.Is(err, os.ErrNotExist) {
			return fmt.Errorf("kimi: remove stale recovery backup: %w", err)
		}
	}
	for index := range staged {
		tmpFile, err := os.CreateTemp(filepath.Dir(staged[index].path), ".agentteam-kimi-recovery-*")
		if err != nil {
			return err
		}
		staged[index].tmp = tmpFile.Name()
		if _, err := tmpFile.Write(contents[index]); err != nil {
			tmpFile.Close()
			return err
		}
		if err := tmpFile.Sync(); err != nil {
			tmpFile.Close()
			return err
		}
		if err := tmpFile.Close(); err != nil {
			return err
		}
	}
	journalCreated := false
	cleanup := func() {
		for _, file := range staged {
			_ = os.Remove(file.tmp)
			if !journalCreated {
				_ = os.Remove(file.backup)
			}
		}
	}
	defer cleanup()
	for index, file := range staged {
		if files[index].originallyMissing {
			if err := os.WriteFile(file.backup, nil, 0o600); err != nil {
				return fmt.Errorf("kimi: mark missing recovery file %s: %w", filepath.Base(file.path), err)
			}
			continue
		}
		if err := os.Link(file.path, file.backup); err != nil {
			return fmt.Errorf("kimi: back up %s: %w", filepath.Base(file.path), err)
		}
	}
	journalPath := filepath.Join(sessionDir, ".agentteam-recovery.json")
	journalFiles := make([]map[string]string, 0, len(staged))
	for _, file := range staged {
		originallyMissing := false
		for _, replacement := range files {
			if replacement.path == file.path {
				originallyMissing = replacement.originallyMissing
				break
			}
		}
		journalFiles = append(journalFiles, map[string]string{
			"path": filepath.Base(file.path), "backup": filepath.Base(file.backup),
			"originally_missing": strconv.FormatBool(originallyMissing),
		})
	}
	journalData, err := json.Marshal(map[string]any{"version": 1, "files": journalFiles})
	if err != nil {
		return err
	}
	if err := writeKimiSyncedFile(journalPath, append(journalData, '\n'), 0o600); err != nil {
		return fmt.Errorf("kimi: write recovery journal: %w", err)
	}
	journalCreated = true
	if err := syncKimiDirectory(sessionDir); err != nil {
		return err
	}
	for _, file := range staged {
		if err := rename(file.tmp, file.path); err != nil {
			rollbackErr := recoverKimiRecoveryTransactionWithRename(sessionDir, rename)
			if rollbackErr == nil {
				journalCreated = false
			}
			return errors.Join(fmt.Errorf("kimi: install restored %s: %w", filepath.Base(file.path), err), rollbackErr)
		}
		if err := syncKimiDirectory(sessionDir); err != nil {
			return err
		}
	}
	if err := os.Remove(journalPath); err != nil {
		return err
	}
	if err := syncKimiDirectory(sessionDir); err != nil {
		return err
	}
	for _, file := range staged {
		if err := os.Remove(file.backup); err != nil && !errors.Is(err, os.ErrNotExist) {
			return err
		}
	}
	journalCreated = false
	return nil
}

func recoverKimiRecoveryTransaction(sessionDir string) error {
	return recoverKimiRecoveryTransactionWithRename(sessionDir, replaceKimiRecoveryFile)
}

func recoverKimiRecoveryTransactionWithRename(
	sessionDir string,
	rename func(string, string) error,
) error {
	journalPath := filepath.Join(sessionDir, ".agentteam-recovery.json")
	data, err := os.ReadFile(journalPath)
	if errors.Is(err, os.ErrNotExist) {
		return nil
	}
	if err != nil {
		return err
	}
	var journal struct {
		Version int `json:"version"`
		Files   []struct {
			Path              string `json:"path"`
			Backup            string `json:"backup"`
			OriginallyMissing string `json:"originally_missing,omitempty"`
		} `json:"files"`
	}
	if err := json.Unmarshal(data, &journal); err != nil || journal.Version != 1 || len(journal.Files) == 0 || len(journal.Files) > 8 {
		if err == nil {
			err = fmt.Errorf("invalid recovery journal")
		}
		return err
	}
	for _, file := range journal.Files {
		if !filepath.IsLocal(file.Path) || !filepath.IsLocal(file.Backup) ||
			filepath.Base(file.Path) != file.Path || filepath.Base(file.Backup) != file.Backup {
			return fmt.Errorf("kimi: unsafe recovery journal path")
		}
		target := filepath.Join(sessionDir, file.Path)
		backup := filepath.Join(sessionDir, file.Backup)
		info, statErr := os.Lstat(backup)
		if statErr != nil {
			return fmt.Errorf("kimi: inspect recovery backup %s: %w", file.Backup, statErr)
		}
		if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
			return fmt.Errorf("kimi: unsafe recovery backup %s", file.Backup)
		}
		originallyMissing, parseErr := strconv.ParseBool(firstNonEmptyString(file.OriginallyMissing, "false"))
		if parseErr != nil {
			return fmt.Errorf("kimi: invalid missing-file recovery marker for %s", file.Path)
		}
		if originallyMissing {
			if err := os.Remove(target); err != nil && !errors.Is(err, os.ErrNotExist) {
				return err
			}
			continue
		}
		restorationPath := backup + ".restoring"
		if err := os.Remove(restorationPath); err != nil && !errors.Is(err, os.ErrNotExist) {
			return fmt.Errorf("kimi: remove stale recovery staging %s: %w", filepath.Base(restorationPath), err)
		}
		if err := os.Link(backup, restorationPath); err != nil {
			return fmt.Errorf("kimi: preserve recovery backup %s: %w", file.Backup, err)
		}
		if err := os.Remove(target); err != nil && !errors.Is(err, os.ErrNotExist) {
			_ = os.Remove(restorationPath)
			return err
		}
		if err := rename(restorationPath, target); err != nil {
			return err
		}
	}
	if err := os.Remove(journalPath); err != nil {
		return err
	}
	if err := syncKimiDirectory(sessionDir); err != nil {
		return err
	}
	for _, file := range journal.Files {
		if err := os.Remove(filepath.Join(sessionDir, file.Backup)); err != nil && !errors.Is(err, os.ErrNotExist) {
			return err
		}
	}
	return nil
}

func writeKimiSyncedFile(path string, data []byte, mode os.FileMode) error {
	file, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, mode)
	if err != nil {
		return err
	}
	if _, err := file.Write(data); err != nil {
		file.Close()
		return err
	}
	if err := file.Sync(); err != nil {
		file.Close()
		return err
	}
	return file.Close()
}

func kimiJSONLBytes(lines []string) []byte {
	if len(lines) == 0 {
		return nil
	}
	return []byte(strings.Join(lines, "\n") + "\n")
}

var _ drivercontract.RecoveryDriver = (*KimiWireDriver)(nil)
var _ drivercontract.DetachedRecoveryDriver = (*KimiWireDriver)(nil)
var _ RuntimeInputReconciler = (*KimiWireDriver)(nil)