AX-Copilot-Codex/src/AxCopilot/Services/Agent/CheckpointTool.cs

using System.IO;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json;

namespace AxCopilot.Services.Agent;

/// <summary>
/// 파일 시스템 체크포인트(스냅샷)를 생성/복원/삭제하는 도구.
/// 작업 폴더의 텍스트 파일을 .ax/checkpoints/ 에 백업하여 undo/rollback을 지원합니다.
/// </summary>
public class CheckpointTool : IAgentTool
{
    public string Name => "checkpoint";

    public string Description =>
        "Create, list, or restore file system checkpoints for undo/rollback. " +
        "Checkpoints capture text files in the working folder as snapshots. " +
        "- action=\"create\": Create a new checkpoint (name optional)\n" +
        "- action=\"list\": List all checkpoints\n" +
        "- action=\"restore\": Restore files from a checkpoint (id or name required, requires user approval)\n" +
        "- action=\"delete\": Delete a checkpoint (id required)";

    public ToolParameterSchema Parameters => new()
    {
        Properties = new()
        {
            ["action"] = new()
            {
                Type = "string",
                Description = "create | list | restore | delete",
                Enum = ["create", "list", "restore", "delete"],
            },
            ["name"] = new()
            {
                Type = "string",
                Description = "Checkpoint name (for create/restore)",
            },
            ["id"] = new()
            {
                Type = "integer",
                Description = "Checkpoint ID (for restore/delete)",
            },
        },
        Required = ["action"],
    };

    private const int MaxCheckpoints = 10;
    private const long MaxFileSize = 5 * 1024 * 1024; // 5MB

    private static readonly HashSet<string> SkipDirs = new(StringComparer.OrdinalIgnoreCase)
    {
        "bin", "obj", "node_modules", ".git", ".vs", ".ax", "packages",
        "Debug", "Release", "TestResults", ".idea", "__pycache__",
    };

    private static readonly HashSet<string> TextExtensions = new(StringComparer.OrdinalIgnoreCase)
    {
        ".cs", ".py", ".js", ".ts", ".java", ".cpp", ".c", ".h", ".hpp",
        ".xml", ".json", ".yaml", ".yml", ".md", ".txt", ".csv", ".html", ".htm",
        ".css", ".sql", ".sh", ".bat", ".ps1", ".config", ".ini", ".xaml",
        ".csproj", ".sln", ".props", ".targets", ".editorconfig", ".gitignore",
        ".tsx", ".jsx", ".vue", ".svelte", ".scss", ".less", ".toml", ".env",
        ".razor", ".proto", ".graphql", ".rs", ".go", ".rb", ".php", ".swift",
    };

    public async Task<ToolResult> ExecuteAsync(JsonElement args, AgentContext context, CancellationToken ct = default)
    {
        if (!args.TryGetProperty("action", out var actionEl))
            return ToolResult.Fail("action이 필요합니다.");
        var action = actionEl.GetString() ?? "";

        if (string.IsNullOrEmpty(context.WorkFolder))
            return ToolResult.Fail("작업 폴더가 설정되지 않았습니다.");

        var checkpointDir = Path.Combine(context.WorkFolder, ".ax", "checkpoints");

        return action switch
        {
            "create" => await CreateCheckpoint(args, context, checkpointDir, ct),
            "list" => ListCheckpoints(checkpointDir),
            "restore" => await RestoreCheckpoint(args, context, checkpointDir, ct),
            "delete" => DeleteCheckpoint(args, checkpointDir),
            _ => ToolResult.Fail($"알 수 없는 액션: {action}. create | list | restore | delete 중 선택하세요."),
        };
    }

    private async Task<ToolResult> CreateCheckpoint(JsonElement args, AgentContext context, string checkpointDir, CancellationToken ct)
    {
        var name = args.TryGetProperty("name", out var n) ? n.GetString() ?? "unnamed" : "unnamed";
        // 이름에서 파일 시스템 비안전 문자 제거
        name = string.Join("_", name.Split(Path.GetInvalidFileNameChars()));

        var timestamp = DateTime.Now.ToString("yyyyMMdd_HHmmss");
        var folderName = $"{timestamp}_{name}";
        var cpDir = Path.Combine(checkpointDir, folderName);

        try
        {
            Directory.CreateDirectory(cpDir);

            // 추적 대상 텍스트 파일 수집
            var files = CollectTextFiles(context.WorkFolder);
            if (files.Count == 0)
                return ToolResult.Fail("체크포인트할 텍스트 파일이 없습니다.");

            var manifest = new CheckpointManifest
            {
                Name = name,
                CreatedAt = DateTime.Now.ToString("yyyy-MM-dd HH:mm:ss"),
                WorkFolder = context.WorkFolder,
                Files = [],
            };

            foreach (var file in files)
            {
                ct.ThrowIfCancellationRequested();

                var relativePath = Path.GetRelativePath(context.WorkFolder, file);
                var destPath = Path.Combine(cpDir, relativePath);
                var destDir = Path.GetDirectoryName(destPath);
                if (!string.IsNullOrEmpty(destDir) && !Directory.Exists(destDir))
                    Directory.CreateDirectory(destDir);

                await CopyFileAsync(file, destPath, ct);

                // SHA256 해시 계산
                var hash = await ComputeHashAsync(file, ct);
                manifest.Files.Add(new FileEntry { Path = relativePath, Hash = hash });
            }

            // manifest.json 저장
            var manifestJson = JsonSerializer.Serialize(manifest, new JsonSerializerOptions
            {
                WriteIndented = true,
                Encoder = System.Text.Encodings.Web.JavaScriptEncoder.UnsafeRelaxedJsonEscaping,
            });
            await TextFileCodec.WriteAllTextAsync(
                Path.Combine(cpDir, "manifest.json"),
                manifestJson,
                TextFileCodec.Utf8NoBom,
                ct);

            // 최대 개수 초과 시 가장 오래된 체크포인트 삭제
            CleanupOldCheckpoints(checkpointDir);

            return ToolResult.Ok($"체크포인트 생성 완료: {folderName}\n파일 수: {manifest.Files.Count}");
        }
        catch (Exception ex)
        {
            return ToolResult.Fail($"체크포인트 생성 오류: {ex.Message}");
        }
    }

    private static ToolResult ListCheckpoints(string checkpointDir)
    {
        if (!Directory.Exists(checkpointDir))
            return ToolResult.Ok("저장된 체크포인트가 없습니다.");

        var dirs = Directory.GetDirectories(checkpointDir)
            .OrderByDescending(d => d)
            .ToList();

        if (dirs.Count == 0)
            return ToolResult.Ok("저장된 체크포인트가 없습니다.");

        var sb = new StringBuilder();
        sb.AppendLine($"체크포인트 {dirs.Count}개:");
        for (var i = 0; i < dirs.Count; i++)
        {
            var dirName = Path.GetFileName(dirs[i]);
            var manifestPath = Path.Combine(dirs[i], "manifest.json");
            var fileCount = 0;
            var createdAt = "";

            if (File.Exists(manifestPath))
            {
                try
                {
                    var json = TextFileCodec.ReadAllText(manifestPath).Text;
                    var manifest = JsonSerializer.Deserialize<CheckpointManifest>(json);
                    fileCount = manifest?.Files.Count ?? 0;
                    createdAt = manifest?.CreatedAt ?? "";
                }
                catch { /* manifest 읽기 실패 무시 */ }
            }

            sb.AppendLine($"  [{i}] {dirName} — 파일 {fileCount}개{(string.IsNullOrEmpty(createdAt) ? "" : $", {createdAt}")}");
        }

        return ToolResult.Ok(sb.ToString());
    }

    private async Task<ToolResult> RestoreCheckpoint(JsonElement args, AgentContext context, string checkpointDir, CancellationToken ct)
    {
        if (!Directory.Exists(checkpointDir))
            return ToolResult.Fail("저장된 체크포인트가 없습니다.");

        var dirs = Directory.GetDirectories(checkpointDir)
            .OrderByDescending(d => d)
            .ToList();

        if (dirs.Count == 0)
            return ToolResult.Fail("저장된 체크포인트가 없습니다.");

        // ID 또는 이름으로 체크포인트 찾기
        string? targetDir = null;

        if (args.TryGetProperty("id", out var idEl))
        {
            var id = idEl.ValueKind == JsonValueKind.Number ? idEl.GetInt32() : int.TryParse(idEl.GetString(), out var parsed) ? parsed : -1;
            if (id >= 0 && id < dirs.Count)
                targetDir = dirs[id];
        }

        if (targetDir == null && args.TryGetProperty("name", out var nameEl))
        {
            var name = nameEl.GetString() ?? "";
            targetDir = dirs.FirstOrDefault(d => Path.GetFileName(d).Contains(name, StringComparison.OrdinalIgnoreCase));
        }

        if (targetDir == null)
            return ToolResult.Fail("체크포인트를 찾을 수 없습니다. id 또는 name을 확인하세요.");

        // 사용자 승인 요청
        if (context.AskPermission != null)
        {
            var approved = await context.AskPermission("checkpoint_restore", $"체크포인트 복원: {Path.GetFileName(targetDir)}");
            if (!approved)
                return ToolResult.Ok("사용자가 복원을 거부했습니다.");
        }

        try
        {
            var manifestPath = Path.Combine(targetDir, "manifest.json");
            if (!File.Exists(manifestPath))
                return ToolResult.Fail("체크포인트 매니페스트를 찾을 수 없습니다.");

            var manifestJson = (await TextFileCodec.ReadAllTextAsync(manifestPath, ct)).Text;
            var manifest = JsonSerializer.Deserialize<CheckpointManifest>(manifestJson);
            if (manifest == null)
                return ToolResult.Fail("매니페스트 파싱 오류");

            var restoredCount = 0;
            foreach (var entry in manifest.Files)
            {
                ct.ThrowIfCancellationRequested();

                var srcPath = Path.Combine(targetDir, entry.Path);
                var destPath = Path.Combine(context.WorkFolder, entry.Path);

                if (!File.Exists(srcPath)) continue;

                var destDir = Path.GetDirectoryName(destPath);
                if (!string.IsNullOrEmpty(destDir) && !Directory.Exists(destDir))
                    Directory.CreateDirectory(destDir);

                await CopyFileAsync(srcPath, destPath, ct);
                restoredCount++;
            }

            return ToolResult.Ok($"체크포인트 복원 완료: {Path.GetFileName(targetDir)}\n복원 파일 수: {restoredCount}/{manifest.Files.Count}");
        }
        catch (Exception ex)
        {
            return ToolResult.Fail($"체크포인트 복원 오류: {ex.Message}");
        }
    }

    private static ToolResult DeleteCheckpoint(JsonElement args, string checkpointDir)
    {
        if (!Directory.Exists(checkpointDir))
            return ToolResult.Fail("저장된 체크포인트가 없습니다.");

        var dirs = Directory.GetDirectories(checkpointDir)
            .OrderByDescending(d => d)
            .ToList();

        if (!args.TryGetProperty("id", out var idEl))
            return ToolResult.Fail("삭제할 체크포인트 id가 필요합니다.");

        var id = idEl.ValueKind == JsonValueKind.Number ? idEl.GetInt32() : int.TryParse(idEl.GetString(), out var parsed) ? parsed : -1;
        if (id < 0 || id >= dirs.Count)
            return ToolResult.Fail($"잘못된 체크포인트 ID: {id}. 0~{dirs.Count - 1} 범위를 사용하세요.");

        try
        {
            var target = dirs[id];
            var name = Path.GetFileName(target);
            Directory.Delete(target, recursive: true);
            return ToolResult.Ok($"체크포인트 삭제됨: {name}");
        }
        catch (Exception ex)
        {
            return ToolResult.Fail($"체크포인트 삭제 오류: {ex.Message}");
        }
    }

    #region Helpers

    private List<string> CollectTextFiles(string workFolder)
    {
        var files = new List<string>();
        CollectFilesRecursive(workFolder, workFolder, files);
        return files;
    }

    private void CollectFilesRecursive(string dir, string rootDir, List<string> files)
    {
        var dirName = Path.GetFileName(dir);
        if (dir != rootDir && SkipDirs.Contains(dirName))
            return;

        try
        {
            foreach (var file in Directory.GetFiles(dir))
            {
                var ext = Path.GetExtension(file);
                if (!TextExtensions.Contains(ext)) continue;

                var fi = new FileInfo(file);
                if (fi.Length > MaxFileSize) continue;

                files.Add(file);
            }

            foreach (var subDir in Directory.GetDirectories(dir))
                CollectFilesRecursive(subDir, rootDir, files);
        }
        catch (UnauthorizedAccessException) { /* 접근 불가 디렉토리 무시 */ }
    }

    private static async Task CopyFileAsync(string src, string dest, CancellationToken ct)
    {
        var buffer = new byte[81920];
        await using var srcStream = new FileStream(src, FileMode.Open, FileAccess.Read, FileShare.Read, buffer.Length, useAsync: true);
        await using var destStream = new FileStream(dest, FileMode.Create, FileAccess.Write, FileShare.None, buffer.Length, useAsync: true);
        int bytesRead;
        while ((bytesRead = await srcStream.ReadAsync(buffer, ct)) > 0)
            await destStream.WriteAsync(buffer.AsMemory(0, bytesRead), ct);
    }

    private static async Task<string> ComputeHashAsync(string filePath, CancellationToken ct)
    {
        using var sha256 = SHA256.Create();
        await using var stream = new FileStream(filePath, FileMode.Open, FileAccess.Read, FileShare.Read, 81920, useAsync: true);
        var hash = await sha256.ComputeHashAsync(stream, ct);
        return Convert.ToHexString(hash).ToLowerInvariant();
    }

    private static void CleanupOldCheckpoints(string checkpointDir)
    {
        if (!Directory.Exists(checkpointDir)) return;
        var dirs = Directory.GetDirectories(checkpointDir)
            .OrderBy(d => d)
            .ToList();

        while (dirs.Count > MaxCheckpoints)
        {
            try
            {
                Directory.Delete(dirs[0], recursive: true);
                dirs.RemoveAt(0);
            }
            catch { break; }
        }
    }

    #endregion

    #region Models

    private class CheckpointManifest
    {
        public string Name { get; set; } = "";
        public string CreatedAt { get; set; } = "";
        public string WorkFolder { get; set; } = "";
        public List<FileEntry> Files { get; set; } = [];
    }

    private class FileEntry
    {
        public string Path { get; set; } = "";
        public string Hash { get; set; } = "";
    }

    #endregion
}