using System.Text.Json;
using Agent.Core.Ollama;
using Agent.Core.Prompts;
using Agent.Core.Sessions;
using Agent.Core.Skills;
using Agent.Core.Tasks;
using Agent.Core.Tools;
using Agent.Planner;

namespace Agent.Core.Runtime;

/// <summary>
/// Builds an initial TaskGraph from LLM structured output (Planning phase).
/// The Planner determines requires_external_effect per-task during Planning.
/// Pure analysis tasks get a TaskGraph with at least 1 Task (RequiresExternalEffect=false).
/// </summary>
public sealed class PlannerService(IOllamaClient ollama, AgentRuntimeSettings settings)
{
    private readonly IOllamaClient _ollama = ollama;
    private readonly AgentRuntimeSettings _settings = settings;

    /// <summary>
    /// Generates an initial TaskGraph for the given request via Ollama structured output.
    /// Returns a TaskGraph stored in the session, keyed by the MCP Task ID.
    /// </summary>
    public async Task<TaskGraph> PlanAsync(Guid taskId, Guid sessionId, AgentSession session, string request, double? temperature, IReadOnlyList<string>? images,
        Action<string, IReadOnlyList<OllamaMessage>, OllamaResult?, Exception?>? observeAttempt, CancellationToken cancellationToken)
    {
        // Build a Planning prompt with skill catalog and sandbox capability summary.
        IReadOnlyList<AgentSkill> selectedSkills = await SelectSkillsAsync(request, temperature, observeAttempt, cancellationToken).ConfigureAwait(false);

        var messages = new List<OllamaMessage>
        {
            new("system", FixedPrompt.Load()),
            new("system", _settings.UserSystemPrompt),
            new("system", $"Respond in locale: {_settings.Locale}."),
            new("system", ComposeSkillCatalog(_settings.Skills)),
            new("system", _settings.SandboxCapabilitySummary),
            new("system", BuildPlanningInstruction()),
            new("user", request, images)
        };

        TaskGraph graph = await RequestValidGraphAsync(messages, temperature, "planning", observeAttempt,
            plan => BuildTaskGraph(plan, request), cancellationToken).ConfigureAwait(false);
        session.TaskGraphs[taskId] = graph;
        return graph;
    }

    /// <summary>
    /// Replans after a policy deny or task failure by adding/replacing tasks in the existing TaskGraph.
    /// Completed tasks are immutable; only incomplete tasks are modified.
    /// </summary>
    public async Task<TaskGraph> ReplanAsync(Guid taskId, Guid sessionId, AgentSession session, string request, double? temperature, string replanReason,
        Action<string, IReadOnlyList<OllamaMessage>, OllamaResult?, Exception?>? observeAttempt, CancellationToken cancellationToken)
    {
        if (!session.TaskGraphs.TryGetValue(taskId, out TaskGraph? existing))
            return await PlanAsync(taskId, sessionId, session, request, temperature, null, observeAttempt, cancellationToken).ConfigureAwait(false);

        var messages = new List<OllamaMessage>
        {
            new("system", FixedPrompt.Load()),
            new("system", _settings.UserSystemPrompt),
            new("system", $"Respond in locale: {_settings.Locale}."),
            new("system", _settings.SandboxCapabilitySummary),
            new("system", BuildReplanInstruction(existing, replanReason)),
            new("user", request)
        };

        TaskGraph rebuilt = await RequestValidGraphAsync(messages, temperature, "replanning", observeAttempt,
            plan => RebuildTaskGraph(existing, plan, request), cancellationToken).ConfigureAwait(false);
        session.TaskGraphs[taskId] = rebuilt;
        return rebuilt;
    }

    private TaskGraph BuildTaskGraph(PlanResult plan, string originalRequest)
    {
        var graph = new TaskGraph(_settings.MaxTasks);
        var idMap = new List<Guid>();

        for (int i = 0; i < plan.Tasks.Count; i++)
        {
            PlannedTask pt = plan.Tasks[i];
            TokenBudgetClass budgetClass = ParseBudgetClass(pt.BudgetClass);
            TokenEstimate estimate = CreateEstimate(pt, budgetClass);

            // Map dependency indices to Guids.
            var dependencies = new HashSet<Guid>();
            foreach (int depIdx in pt.Dependencies)
            {
                if (depIdx < 0 || depIdx >= idMap.Count)
                    throw new PlannerValidationException($"Task {i} contains invalid dependency index {depIdx}.");
                dependencies.Add(idMap[depIdx]);
            }

            string description = string.IsNullOrWhiteSpace(pt.Description) ? originalRequest : pt.Description;
            PlannerTask task = graph.AddTask(description, budgetClass, estimate, dependencies, pt.RequiresExternalEffect,
                requiredTools: ValidateRequiredTools(pt));
            idMap.Add(task.Id);
        }

        if (graph.Count == 0) throw new PlannerValidationException("Structured plan did not contain an executable Task.");

        return graph;
    }

    private TaskGraph RebuildTaskGraph(TaskGraph existing, PlanResult plan, string originalRequest)
    {
        PlannerTask[] completed = existing.Tasks.Where(static t => t.State == PlannerTaskState.Completed).ToArray();
        var completedByDescription = completed.GroupBy(static t => t.Description, StringComparer.Ordinal)
            .ToDictionary(static group => group.Key, static group => group.First(), StringComparer.Ordinal);
        var graph = new TaskGraph(_settings.MaxTasks);
        foreach (PlannerTask item in completed)
        {
            Guid[] dependencies = item.Dependencies.Where(id => completed.Any(candidate => candidate.Id == id)).ToArray();
            PlannerTask copy = graph.AddTask(item.Description, item.BudgetClass, item.Estimate, dependencies, item.RequiresExternalEffect, item.Id, item.RequiredTools);
            graph.AddUsage(copy.Id, item.ActualInputTokens, item.ActualOutputTokens);
            graph.SetState(copy.Id, PlannerTaskState.Completed);
        }
        var idMap = new List<Guid>();
        var reusedCompleted = new HashSet<Guid>();

        for (int i = 0; i < plan.Tasks.Count; i++)
        {
            PlannedTask pt = plan.Tasks[i];

            // Skip tasks that are already completed in the existing graph.
            if (completedByDescription.TryGetValue(pt.Description, out PlannerTask? completedTask) && reusedCompleted.Add(completedTask.Id))
            {
                idMap.Add(completedTask.Id);
                continue;
            }

            TokenBudgetClass bc = ParseBudgetClass(pt.BudgetClass);
            TokenEstimate est = CreateEstimate(pt, bc);
            var dependencies = new HashSet<Guid>();
            foreach (int depIdx in pt.Dependencies)
            {
                if (depIdx < 0 || depIdx >= idMap.Count)
                    throw new PlannerValidationException($"Task {i} contains invalid dependency index {depIdx}.");
                dependencies.Add(idMap[depIdx]);
            }

            string description = string.IsNullOrWhiteSpace(pt.Description) ? originalRequest : pt.Description;
            PlannerTask task = graph.AddTask(description, bc, est, dependencies, pt.RequiresExternalEffect,
                requiredTools: ValidateRequiredTools(pt));
            idMap.Add(task.Id);
        }

        if (graph.Tasks.All(static task => task.State == PlannerTaskState.Completed))
            throw new PlannerValidationException("Replan did not produce an incomplete executable Task.");
        return graph;
    }

    private static TokenBudgetClass ParseBudgetClass(string value)
    {
        try { return TokenBudgetClassExtensions.Parse(value); }
        catch { return TokenBudgetClass.K4; }
    }

    private static TokenEstimate CreateEstimate(PlannedTask pt, TokenBudgetClass budgetClass)
    {
        long o = Math.Max(1, pt.OptimisticTokens);
        long m = Math.Max(o, pt.MostLikelyTokens);
        long p = Math.Max(m, pt.PessimisticTokens);
        long max = budgetClass.MaxTokens();
        if (p > max) p = max;
        if (m > p) m = p;
        if (o > m) o = m;
        return new TokenEstimate(o, m, p);
    }

    private static IReadOnlySet<string> ValidateRequiredTools(PlannedTask task)
    {
        HashSet<string> tools = task.RequiredTools.ToHashSet(StringComparer.Ordinal);
        string[] allowed = ["file_search", "file_read", "file_edit", "process_run", "task_update", "dir"];
        if (tools.Any(tool => !allowed.Contains(tool, StringComparer.Ordinal)))
            throw new PlannerValidationException("required_tools contains an unknown Internal Tool.");
        if (task.RequiresExternalEffect && tools.Count == 0)
            throw new PlannerValidationException("An external-effect Task must provide required_tools.");
        if (!task.RequiresExternalEffect && tools.Any(tool => tool is "file_edit" or "process_run"))
            throw new PlannerValidationException("A non-effect Task cannot require an effect-only Tool.");
        return tools;
    }

    private static PlanResult? ParsePlan(string content)
    {
        try
        {
            return JsonSerializer.Deserialize<PlanResult>(StructuredJson.Unwrap(content), new JsonSerializerOptions
            {
                PropertyNameCaseInsensitive = true
            });
        }
        catch
        {
            return null;
        }
    }

    private async Task<IReadOnlyList<AgentSkill>> SelectSkillsAsync(string request, double? temperature,
        Action<string, IReadOnlyList<OllamaMessage>, OllamaResult?, Exception?>? observeAttempt, CancellationToken cancellationToken)
    {
        if (_settings.Skills.Count == 0) return [];

        string catalog = string.Join("\n", _settings.Skills.Select(static s => $"- {s.Name}: {s.Description}"));
        var selectionMessages = new List<OllamaMessage>
        {
            new("system", $"You are a skill classifier, not a task executor. Do not solve the request, inspect files, write code, or call Tools. Available skills:\n{catalog}\n\nReturn only one JSON object with a 'selected_skills' array containing exact skill names. Select only relevant skills. If none are relevant, return an empty array."),
            new("user", $"Classify skills for this request only:\n{request}")
        };

        string schema = """{"type":"object","properties":{"selected_skills":{"type":"array","items":{"type":"string"}}},"required":["selected_skills"]}""";
        using JsonDocument formatDoc = JsonDocument.Parse(schema);

        IReadOnlyList<OllamaMessage> attemptMessages = selectionMessages;
        for (int attempt = 0; attempt < 2; attempt++)
        {
            OllamaResult? result = null;
            try
            {
                bool? thinkOverride = attempt == 0 ? null : false;
                result = await _ollama.ChatStructuredAsync(attemptMessages, temperature, formatDoc.RootElement.Clone(), thinkOverride, cancellationToken).ConfigureAwait(false);
                if (!TryParseSelectedSkills(result.Content, out IReadOnlyList<AgentSkill> selected))
                    throw new PlannerValidationException("Ollama returned an invalid skill selection object.");
                observeAttempt?.Invoke("skill_selection", attemptMessages, result, null);
                return selected;
            }
            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) { throw; }
            catch (Exception ex) when (ex is not (OutOfMemoryException or StackOverflowException or AccessViolationException))
            {
                observeAttempt?.Invoke("skill_selection", attemptMessages, result, ex);
                if (attempt == 0)
                {
                    attemptMessages = selectionMessages.Concat([
                        new OllamaMessage("system", "The previous response was invalid. Return only {\"selected_skills\":[\"name\"]} or {\"selected_skills\":[]}, with no Markdown, code, Tool call, or commentary.")
                    ]).ToArray();
                }
            }
        }
        return [];
    }

    private async Task<TaskGraph> RequestValidGraphAsync(IReadOnlyList<OllamaMessage> messages, double? temperature, string phase,
        Action<string, IReadOnlyList<OllamaMessage>, OllamaResult?, Exception?>? observeAttempt, Func<PlanResult, TaskGraph> buildGraph, CancellationToken cancellationToken)
    {
        Exception? lastFailure = null;
        IReadOnlyList<OllamaMessage> attemptMessages = messages;
        for (int attempt = 0; attempt < 2; attempt++)
        {
            OllamaResult? result = null;
            try
            {
                bool? thinkOverride = attempt == 0 ? null : false;
                result = await _ollama.ChatStructuredAsync(attemptMessages, temperature, PlanSchema.Element, thinkOverride, cancellationToken).ConfigureAwait(false);
                PlanResult? plan = ParsePlan(result.Content);
                if (plan is null || plan.Tasks.Count == 0)
                    throw new PlannerValidationException("Ollama returned an empty or invalid structured plan.");
                TaskGraph graph = buildGraph(plan);
                observeAttempt?.Invoke(phase, attemptMessages, result, null);
                return graph;
            }
            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) { throw; }
            catch (Exception ex) when (ex is not (OutOfMemoryException or StackOverflowException or AccessViolationException))
            {
                observeAttempt?.Invoke(phase, attemptMessages, result, ex);
                lastFailure = ex;
                if (attempt == 0)
                {
                    attemptMessages = messages.Concat([
                        new OllamaMessage("system",
                            "The previous response was not a valid TaskGraph JSON object. Do not call or execute any Tool while planning. " +
                            "Return only one JSON object that matches the supplied schema, beginning with {\"tasks\":[ and containing no Markdown or commentary.")
                    ]).ToArray();
                }
            }
        }
        throw new PlannerValidationException($"Structured {phase} failed after think=false retry: {lastFailure?.Message}");
    }

    private bool TryParseSelectedSkills(string content, out IReadOnlyList<AgentSkill> selected)
    {
        selected = [];
        try
        {
            using JsonDocument doc = JsonDocument.Parse(StructuredJson.Unwrap(content));
            if (doc.RootElement.ValueKind != JsonValueKind.Object ||
                !doc.RootElement.TryGetProperty("selected_skills", out JsonElement skillsElement) ||
                skillsElement.ValueKind != JsonValueKind.Array) return false;
            var selectedNames = new HashSet<string>(StringComparer.Ordinal);
            foreach (JsonElement item in skillsElement.EnumerateArray())
            {
                if (item.ValueKind != JsonValueKind.String) return false;
                string? name = item.GetString();
                if (!string.IsNullOrWhiteSpace(name)) selectedNames.Add(name);
            }
            var valid = new List<AgentSkill>();
            foreach (AgentSkill skill in _settings.Skills)
            {
                if (selectedNames.Contains(skill.Name))
                {
                    valid.Add(skill);
                    if (valid.Count >= 10) break;
                }
            }
            selected = valid;
            return true;
        }
        catch
        {
            return false;
        }
    }

    private static string ComposeSkillCatalog(IReadOnlyList<AgentSkill> skills)
    {
        if (skills.Count == 0) return "No skills available.";
        var lines = new List<string> { "Available skills:" };
        foreach (AgentSkill skill in skills)
            lines.Add($"- {skill.Name}: {skill.Description}");
        return string.Join("\n", lines);
    }

    private static string BuildPlanningInstruction() =>
        """
        You are a Task Planner. Decompose the user request into a TaskGraph of 1-32 tasks.
        Planning is read-only: do not call or execute Tools. Return only one JSON object matching the supplied schema, with no Markdown or commentary.
        Each task must have:
        - description: a clear, self-contained sub-task description
        - dependencies: array of task indices (0-based) that must complete before this task
        - budget_class: one of 1K, 2K, 4K, 8K, 16K, 32K, 64K, 128K (token budget for this task)
        - optimistic_tokens, most_likely_tokens, pessimistic_tokens: PERT estimates (0 < O <= M <= P <= budget)
        - requires_external_effect: true if this task requires a verifiable external effect (file edit, process run, etc.)
        - required_tools: the exact Internal Tool names whose successful result is required to complete this task; use only file_search, file_read, file_edit, process_run, task_update, or dir; use [] for a pure response task
        
        For pure analysis or informational tasks, requires_external_effect should be false.
        For tasks that modify files, run processes, or produce verifiable side effects, requires_external_effect should be true.
        Read-only operations (file_read, file_search, and dir list) do not create an external effect and must use requires_external_effect=false.
        A read/search/list task must name file_read, file_search, or dir in required_tools even though requires_external_effect is false. A pure response task uses an empty array. When requires_external_effect is true, required_tools must include the effect Tool that proves this task's own expected effect.
        File creation and file modification use file_edit. There is no file_write Tool. Directory creation and removal use dir.
        Preserve every user-supplied absolute path, executable, cwd, argument array, operation name, and literal value exactly in the task that uses it; never abbreviate these values into a shell-like command.
        Keep each task limited to its own described operation. Do not put later dependent operations into an earlier task.
        For coding work whose edits must be verified, after any read-only inspection create one implementation-and-verification Task requiring both file_edit and process_run. Do not split the initial edit, test edit, first verification, or failure correction into separate Tasks. This lets the same Task make the required edits, stop after an initially successful verification, or correct and rerun after failure. A required Tool with an unsuccessful result cannot complete a Task or defer its correction to a dependent Task.
        A single task is valid for simple requests.
        """;

    private static string BuildReplanInstruction(TaskGraph existing, string reason)
    {
        var sb = new System.Text.StringBuilder();
        sb.AppendLine("You are a Task Planner performing replan. The previous plan encountered an issue.");
        sb.AppendLine($"Replan reason: {reason}");
        sb.AppendLine("Current task states:");
        foreach (PlannerTask t in existing.Tasks)
        {
            sb.AppendLine($"  - [{t.State}] {t.Description} (requires_external_effect={t.RequiresExternalEffect})");
        }
        sb.AppendLine();
        sb.AppendLine("Completed tasks must be preserved. Modify only incomplete/failed tasks.");
        sb.AppendLine("Return a complete revised plan including completed tasks.");
        return sb.ToString();
    }
}