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

using System.Text;
using System.Text.Json;
using AxCopilot.Models;
using AxCopilot.Services;

namespace AxCopilot.Services.Agent;

/// <summary>
/// Spawns a background sub-agent that runs an isolated read-only agent loop.
/// The parent agent can later collect all finished results with wait_agents.
/// </summary>
public class SubAgentTool : IAgentTool
{
    public static event Action<SubAgentStatusEvent>? StatusChanged;

    public string Name => "spawn_agent";

    public string Description =>
        "Create a read-only sub-agent for bounded parallel research or codebase analysis.\n" +
        "Use this when a side task can run independently while the main agent continues.\n" +
        "The sub-agent can inspect files, search code, review diffs, and summarize findings.\n" +
        "Collect results later with wait_agents.";

    public ToolParameterSchema Parameters => new()
    {
        Properties = new()
        {
            ["task"] = new ToolProperty
            {
                Type = "string",
                Description = "The self-contained task for the sub-agent."
            },
            ["id"] = new ToolProperty
            {
                Type = "string",
                Description = "A unique sub-agent identifier used by wait_agents."
            },
            ["profile"] = new ToolProperty
            {
                Type = "string",
                Description = "Execution profile: researcher (default, read-only), coder (can edit/build), writer (doc creation), reviewer (code review), planner (task decomposition)."
            },
        },
        Required = new() { "task", "id" }
    };

    private static readonly Dictionary<string, SubAgentTask> _activeTasks = new();
    private static readonly object _lock = new();

    public Task<ToolResult> ExecuteAsync(JsonElement args, AgentContext context, CancellationToken ct = default)
    {
        var task = args.SafeTryGetProperty("task", out var t) ? t.SafeGetString() ?? "" : "";
        var id = args.SafeTryGetProperty("id", out var i) ? i.SafeGetString() ?? "" : "";
        var profileName = args.SafeTryGetProperty("profile", out var p) ? p.SafeGetString() : null;

        if (string.IsNullOrWhiteSpace(task) || string.IsNullOrWhiteSpace(id))
            return Task.FromResult(ToolResult.Fail("task and id are required."));

        CleanupStale();

        var app = System.Windows.Application.Current as App;
        var maxAgents = app?.SettingsService?.Settings.Llm.MaxSubAgents ?? 3;

        lock (_lock)
        {
            if (_activeTasks.ContainsKey(id))
                return Task.FromResult(ToolResult.Fail($"Sub-agent id already exists: {id}"));

            var running = _activeTasks.Values.Count(x => x.CompletedAt == null);
            if (running >= maxAgents)
                return Task.FromResult(ToolResult.Fail($"Maximum concurrent sub-agents reached ({maxAgents})."));
        }

        var subTask = new SubAgentTask
        {
            Id = id,
            Task = task,
            StartedAt = DateTime.Now,
        };

        // P2: 부모 취소 토큰 연동 — 부모 에이전트 중지 시 자식도 즉시 취소
        var cts = CancellationTokenSource.CreateLinkedTokenSource(ct);
        subTask.Cts = cts;

        subTask.RunTask = System.Threading.Tasks.Task.Run(async () =>
        {
            try
            {
                var result = await RunSubAgentAsync(id, task, context, profileName, cts.Token).ConfigureAwait(false);
                subTask.Result = result;
                subTask.Success = true;
                NotifyStatus(new SubAgentStatusEvent
                {
                    Id = id,
                    Task = task,
                    Status = SubAgentRunStatus.Completed,
                    Summary = $"Sub-agent '{id}' completed.",
                    Result = result,
                    Timestamp = DateTime.Now,
                });
            }
            catch (OperationCanceledException)
            {
                subTask.Result = "Cancelled: parent agent was stopped.";
                subTask.Success = false;
                NotifyStatus(new SubAgentStatusEvent
                {
                    Id = id,
                    Task = task,
                    Status = SubAgentRunStatus.Failed,
                    Summary = $"Sub-agent '{id}' cancelled.",
                    Result = subTask.Result,
                    Timestamp = DateTime.Now,
                });
            }
            catch (Exception ex)
            {
                subTask.Result = $"Error: {ex.Message}";
                subTask.Success = false;
                NotifyStatus(new SubAgentStatusEvent
                {
                    Id = id,
                    Task = task,
                    Status = SubAgentRunStatus.Failed,
                    Summary = $"Sub-agent '{id}' failed: {ex.Message}",
                    Result = subTask.Result,
                    Timestamp = DateTime.Now,
                });
            }
            finally
            {
                subTask.CompletedAt = DateTime.Now;
                cts.Dispose();
            }
        }, cts.Token);

        lock (_lock)
            _activeTasks[id] = subTask;

        NotifyStatus(new SubAgentStatusEvent
        {
            Id = id,
            Task = task,
            Status = SubAgentRunStatus.Started,
            Summary = $"Sub-agent '{id}' started.",
            Timestamp = DateTime.Now,
        });

        return Task.FromResult(ToolResult.Ok(
            $"Sub-agent '{id}' started.\nTask: {task}\nUse wait_agents later to collect the result."));
    }

    private static async Task<string> RunSubAgentAsync(string id, string task, AgentContext parentContext, string? profileName, CancellationToken ct)
    {
        var profile = SubAgentProfileCatalog.Get(profileName);
        var settings = CreateSubAgentSettings(parentContext, profile);
        using var llm = new LlmService(settings);
        // P2: 프로파일별 temperature override
        if (profile.TemperatureOverride.HasValue)
            llm.PushInferenceOverride(temperature: profile.TemperatureOverride.Value);
        using var tools = await CreateSubAgentRegistryAsync(settings, profile).ConfigureAwait(false);

        var loop = new AgentLoopService(llm, tools, settings)
        {
            ActiveTab = parentContext.ActiveTab,
        };

        var messages = new List<ChatMessage>
        {
            new()
            {
                Role = "system",
                Content = BuildSubAgentSystemPrompt(task, parentContext, profile),
            },
            new()
            {
                Role = "user",
                Content = task,
            }
        };

        var finalText = await loop.RunAsync(messages, ct).ConfigureAwait(false);

        var eventSummary = SummarizeEvents(loop.Events);
        var sb = new StringBuilder();
        sb.AppendLine($"[Sub-agent {id}]");
        sb.AppendLine($"Task: {task}");
        if (!string.IsNullOrWhiteSpace(parentContext.WorkFolder))
            sb.AppendLine($"Work folder: {parentContext.WorkFolder}");
        if (!string.IsNullOrWhiteSpace(eventSummary))
        {
            sb.AppendLine();
            sb.AppendLine("Observed work:");
            sb.AppendLine(eventSummary);
        }
        sb.AppendLine();
        sb.AppendLine("Result:");
        sb.AppendLine(string.IsNullOrWhiteSpace(finalText) ? "(empty)" : finalText.Trim());
        return sb.ToString().TrimEnd();
    }

    private static SettingsService CreateSubAgentSettings(AgentContext parentContext, SubAgentProfile profile)
    {
        var settings = new SettingsService();
        settings.Load();

        var llm = settings.Settings.Llm;
        llm.WorkFolder = parentContext.WorkFolder;
        llm.FilePermission = profile.FilePermission;
        llm.AgentHooks = new();
        llm.ToolPermissions = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
        llm.DisabledTools = profile.DisabledToolNames.ToList();

        return settings;
    }

    /// <summary>도구 이름 → 팩토리 매핑 (인스턴스를 필요한 것만 생성).</summary>
    private static readonly Dictionary<string, Func<IAgentTool>> ToolFactories =
        new(StringComparer.OrdinalIgnoreCase)
    {
        ["file_read"] = () => new FileReadTool(),
        ["glob"] = () => new GlobTool(),
        ["grep"] = () => new GrepTool(),
        ["folder_map"] = () => new FolderMapTool(),
        ["document_read"] = () => new DocumentReaderTool(),
        ["dev_env_detect"] = () => new DevEnvDetectTool(),
        ["git_tool"] = () => new GitTool(),
        ["lsp_code_intel"] = () => new LspTool(),
        ["code_search"] = () => new CodeSearchTool(),
        ["code_review"] = () => new CodeReviewTool(),
        ["project_rule"] = () => new ProjectRuleTool(),
        ["skill_manager"] = () => new SkillManagerTool(),
        ["json_tool"] = () => new JsonTool(),
        ["regex_tool"] = () => new RegexTool(),
        ["diff_tool"] = () => new DiffTool(),
        ["base64_tool"] = () => new Base64Tool(),
        ["hash_tool"] = () => new HashTool(),
        ["datetime_tool"] = () => new DateTimeTool(),
        ["math_tool"] = () => new MathTool(),
        ["xml_tool"] = () => new XmlTool(),
        ["multi_read"] = () => new MultiReadTool(),
        ["file_info"] = () => new FileInfoTool(),
        ["document_review"] = () => new DocumentReviewTool(),
        // coder 프로파일용
        ["file_write"] = () => new FileWriteTool(),
        ["file_edit"] = () => new FileEditTool(),
        ["build_run"] = () => new BuildRunTool(),
        ["process"] = () => new ProcessTool(),
        ["test_loop"] = () => new TestLoopTool(),
        ["snippet_runner"] = () => new SnippetRunnerTool(),
        // writer 프로파일용
        ["html_create"] = () => new HtmlSkill(),
        ["docx_create"] = () => new DocxSkill(),
        ["markdown_create"] = () => new MarkdownSkill(),
        ["csv_create"] = () => new CsvSkill(),
        ["excel_create"] = () => new ExcelSkill(),
        ["pptx_create"] = () => new PptxSkill(),
        ["document_plan"] = () => new DocumentPlannerTool(),
    };

    private static async Task<ToolRegistry> CreateSubAgentRegistryAsync(SettingsService settings, SubAgentProfile profile)
    {
        var registry = new ToolRegistry();

        // 필요한 도구만 인스턴스 생성 (기존: 전체 63개 생성 후 필터 → 개선: 필요한 것만 팩토리 호출)
        foreach (var name in profile.EnabledToolNames)
        {
            if (ToolFactories.TryGetValue(name, out var factory))
                registry.Register(factory());
        }

        await registry.RegisterMcpToolsAsync(settings.Settings.Llm.McpServers).ConfigureAwait(false);
        return registry;
    }

    private static string BuildSubAgentSystemPrompt(string task, AgentContext parentContext, SubAgentProfile profile)
    {
        var sb = new StringBuilder();

        // P2: 프로파일별 시스템 프롬프트 접두사 사용
        sb.AppendLine(profile.SystemPromptPrefix);

        if (!string.IsNullOrWhiteSpace(parentContext.WorkFolder))
            sb.AppendLine($"Current work folder: {parentContext.WorkFolder}");
        if (!string.IsNullOrWhiteSpace(parentContext.ActiveTab))
            sb.AppendLine($"Current tab: {parentContext.ActiveTab}");

        // P4: 워크스페이스 컨텍스트 자동 주입
        var wsContext = WorkspaceContextGenerator.LoadContext(parentContext.WorkFolder);
        if (!string.IsNullOrWhiteSpace(wsContext))
        {
            sb.AppendLine();
            sb.AppendLine("Workspace context:");
            sb.AppendLine(wsContext.Length > 2000 ? wsContext[..2000] + "\n...(truncated)" : wsContext);
        }

        sb.AppendLine();

        // 프로파일별 작업 규칙
        switch (profile.Name)
        {
            case "coder":
                sb.AppendLine("Coding rules:");
                sb.AppendLine("1. Read the relevant files first to understand existing patterns.");
                sb.AppendLine("2. Make the minimal correct change — do not refactor unrelated code.");
                sb.AppendLine("3. After editing, verify with build_run or test_loop.");
                sb.AppendLine("4. If the build fails, fix the issue immediately.");
                sb.AppendLine("5. Report what was changed and the verification result.");
                break;

            case "writer":
                sb.AppendLine("Document creation rules:");
                sb.AppendLine("1. Inspect existing documents or source files for context.");
                sb.AppendLine("2. Produce well-structured, complete documents.");
                sb.AppendLine("3. Use appropriate formatting for the target format.");
                sb.AppendLine("4. Verify file was created successfully.");
                break;

            case "reviewer":
                sb.AppendLine("Review rules:");
                sb.AppendLine("1. Start by reading the directly relevant files.");
                sb.AppendLine("2. Rate each finding P0 (critical) through P3 (minor).");
                sb.AppendLine("3. Prioritize concrete defects, regressions, and missing tests.");
                sb.AppendLine("4. Cite exact file paths and line ranges as evidence.");
                sb.AppendLine("5. Do not suggest edits — only report findings.");
                break;

            case "planner":
                sb.AppendLine("Planning rules:");
                sb.AppendLine("1. Inspect the codebase to understand the current architecture.");
                sb.AppendLine("2. Decompose the task into ordered steps with clear dependencies.");
                sb.AppendLine("3. Identify the minimum file set for each step.");
                sb.AppendLine("4. Highlight the primary risk for each step.");
                sb.AppendLine("5. Suggest a validation strategy.");
                break;

            default: // researcher
                sb.AppendLine("Investigation rules:");
                sb.AppendLine("1. Start by reading the directly relevant files, not by summarizing from memory.");
                sb.AppendLine("2. If the task mentions a code path, type, method, or feature, use grep/glob to find references and callers.");
                sb.AppendLine("3. If the task is about bugs, trace likely cause -> affected files -> validation evidence.");
                sb.AppendLine("4. If the task is about implementation planning, identify the minimum file set and the main risk.");
                sb.AppendLine("5. If the task is about review, prioritize concrete defects, regressions, and missing tests.");
                break;
        }

        var workflowHints = BuildSubAgentWorkflowHints(task, parentContext.ActiveTab);
        if (!string.IsNullOrWhiteSpace(workflowHints))
        {
            sb.AppendLine();
            sb.AppendLine("Task-specific guidance:");
            sb.AppendLine(workflowHints);
        }
        sb.AppendLine("Final answer format:");
        sb.AppendLine("1. Short conclusion");
        sb.AppendLine("2. Files checked");
        sb.AppendLine("3. Key evidence");
        sb.AppendLine("4. Recommended next action for the main agent");
        sb.AppendLine("5. Risks or unknowns");
        return sb.ToString().TrimEnd();
    }

    private static string BuildSubAgentWorkflowHints(string task, string? activeTab)
    {
        var normalizedTask = task ?? "";
        var sb = new StringBuilder();

        if (normalizedTask.Contains("review", StringComparison.OrdinalIgnoreCase)
            || normalizedTask.Contains("검토", StringComparison.OrdinalIgnoreCase)
            || normalizedTask.Contains("리뷰", StringComparison.OrdinalIgnoreCase))
        {
            sb.AppendLine("- Review tasks should name the concrete issue first, then cite the supporting file evidence.");
            sb.AppendLine("- Mention missing or weak tests when behavior could regress.");
        }

        if (normalizedTask.Contains("bug", StringComparison.OrdinalIgnoreCase)
            || normalizedTask.Contains("error", StringComparison.OrdinalIgnoreCase)
            || normalizedTask.Contains("실패", StringComparison.OrdinalIgnoreCase)
            || normalizedTask.Contains("오류", StringComparison.OrdinalIgnoreCase))
        {
            sb.AppendLine("- Bug investigations should identify the most likely root cause and the exact files that support that conclusion.");
            sb.AppendLine("- Suggest the smallest safe fix path for the main agent.");
        }

        if (string.Equals(activeTab, "Code", StringComparison.OrdinalIgnoreCase))
        {
            sb.AppendLine("- For code tasks, always mention impacted callers/references if you found any.");
            sb.AppendLine("- Call out related tests or note explicitly when tests were not found.");
        }

        if (normalizedTask.Contains("plan", StringComparison.OrdinalIgnoreCase)
            || normalizedTask.Contains("설계", StringComparison.OrdinalIgnoreCase)
            || normalizedTask.Contains("계획", StringComparison.OrdinalIgnoreCase))
        {
            sb.AppendLine("- Planning tasks should identify the minimum file set, order of work, and the primary validation step.");
        }

        return sb.ToString().TrimEnd();
    }

    private static string SummarizeEvents(IEnumerable<AgentEvent> events)
    {
        var lines = events
            .Where(e => e.Type is AgentEventType.ToolCall or AgentEventType.ToolResult or AgentEventType.StepStart)
            .TakeLast(12)
            .Select(e =>
            {
                var label = e.Type switch
                {
                    AgentEventType.ToolCall => $"tool:{e.ToolName}",
                    AgentEventType.ToolResult => $"result:{e.ToolName}",
                    AgentEventType.StepStart => "step",
                    _ => e.Type.ToString().ToLowerInvariant()
                };
                var summary = string.IsNullOrWhiteSpace(e.Summary) ? "" : $" - {e.Summary.Trim()}";
                return $"- {label}{summary}";
            })
            .ToList();

        return lines.Count == 0 ? "" : string.Join(Environment.NewLine, lines);
    }

    public static IReadOnlyDictionary<string, SubAgentTask> ActiveTasks
    {
        get
        {
            lock (_lock)
                return new Dictionary<string, SubAgentTask>(_activeTasks);
        }
    }

    public static async Task<string> WaitAsync(
        IEnumerable<string>? ids = null,
        bool completedOnly = false,
        CancellationToken ct = default)
    {
        List<SubAgentTask> tasks;
        lock (_lock)
            tasks = _activeTasks.Values.ToList();

        var requestedIds = ids?
            .Where(x => !string.IsNullOrWhiteSpace(x))
            .Select(x => x.Trim())
            .Distinct(StringComparer.OrdinalIgnoreCase)
            .ToList();

        if (requestedIds is { Count: > 0 })
            tasks = tasks
                .Where(t => requestedIds.Contains(t.Id, StringComparer.OrdinalIgnoreCase))
                .ToList();

        if (tasks.Count == 0)
        {
            if (requestedIds is { Count: > 0 })
                return $"No matching sub-agents found for: {string.Join(", ", requestedIds)}";
            return "No active sub-agents.";
        }

        if (!completedOnly)
            await Task.WhenAll(tasks.Where(t => t.RunTask != null).Select(t => t.RunTask!)).WaitAsync(ct);
        else
            tasks = tasks.Where(t => t.CompletedAt != null).ToList();

        if (tasks.Count == 0)
            return "No requested sub-agents have completed yet.";

        var sb = new StringBuilder();
        sb.AppendLine($"Collected {tasks.Count} sub-agent result(s):");
        foreach (var task in tasks.OrderBy(t => t.StartedAt))
        {
            var status = task.Success ? "OK" : "FAIL";
            var duration = task.CompletedAt.HasValue
                ? $"{(task.CompletedAt.Value - task.StartedAt).TotalSeconds:F1}s"
                : "running";
            sb.AppendLine();
            sb.AppendLine($"--- [{status}] {task.Id} ({duration}) ---");
            sb.AppendLine(task.Result ?? "(no result)");
        }

        lock (_lock)
        {
            foreach (var task in tasks)
                _activeTasks.Remove(task.Id);
        }

        return sb.ToString().TrimEnd();
    }

    public static void CleanupStale()
    {
        lock (_lock)
        {
            var stale = _activeTasks
                .Where(kv => kv.Value.CompletedAt.HasValue &&
                             (DateTime.Now - kv.Value.CompletedAt.Value).TotalMinutes > 10)
                .Select(kv => kv.Key)
                .ToList();

            foreach (var key in stale)
                _activeTasks.Remove(key);
        }
    }

    private static void NotifyStatus(SubAgentStatusEvent evt)
    {
        try { StatusChanged?.Invoke(evt); } catch { }
    }
}

public class WaitAgentsTool : IAgentTool
{
    public string Name => "wait_agents";

    public string Description =>
        "Wait for sub-agents and collect their results.\n" +
        "You may wait for all sub-agents or only specific ids.\n" +
        "Use completed_only=true to collect only already finished sub-agents without blocking.";

    public ToolParameterSchema Parameters => new()
    {
        Properties = new()
        {
            ["ids"] = new ToolProperty
            {
                Type = "array",
                Description = "Optional list of sub-agent ids to collect. Omit to collect all.",
                Items = new ToolProperty { Type = "string", Description = "Sub-agent id" }
            },
            ["completed_only"] = new ToolProperty
            {
                Type = "boolean",
                Description = "If true, collect only already completed sub-agents and do not wait."
            },
        },
        Required = new()
    };

    public async Task<ToolResult> ExecuteAsync(JsonElement args, AgentContext context, CancellationToken ct = default)
    {
        List<string>? ids = null;
        if (args.SafeTryGetProperty("ids", out var idsEl) && idsEl.ValueKind == JsonValueKind.Array)
        {
            ids = idsEl.EnumerateArray()
                .Where(x => x.ValueKind == JsonValueKind.String)
                .Select(x => x.SafeGetString() ?? "")
                .Where(x => !string.IsNullOrWhiteSpace(x))
                .ToList();
        }

        var completedOnly = args.SafeTryGetProperty("completed_only", out var completedEl) &&
                            completedEl.ValueKind == JsonValueKind.True;

        var result = await SubAgentTool.WaitAsync(ids, completedOnly, ct).ConfigureAwait(false);
        return ToolResult.Ok(result);
    }
}

public class SubAgentTask
{
    public string Id { get; init; } = "";
    public string Task { get; init; } = "";
    public DateTime StartedAt { get; init; }
    public DateTime? CompletedAt { get; set; }
    public bool Success { get; set; }
    public string? Result { get; set; }
    public Task? RunTask { get; set; }
    public CancellationTokenSource? Cts { get; set; }
}

public enum SubAgentRunStatus
{
    Started,
    Completed,
    Failed
}

public class SubAgentStatusEvent
{
    public string Id { get; init; } = "";
    public string Task { get; init; } = "";
    public SubAgentRunStatus Status { get; init; }
    public string Summary { get; init; } = "";
    public string? Result { get; init; }
    public DateTime Timestamp { get; init; } = DateTime.Now;
}