using System.ComponentModel;
using Agent.Core.Localization;
using Agent.Core.Tasks;
using ModelContextProtocol.Server;
using ModelContextProtocol.Protocol;
using ModelContextProtocol.Extensions.Tasks;
using System.Text.Json;
using System.Text.Json.Nodes;

namespace Agent.Mcp;

[McpServerToolType]
public sealed class AgentTools(TaskRegistry registry, LocaleCatalog locale, VisionLimits visionLimits, PersistentMcpTaskStore taskStore)
{
    [McpServerTool(Name = "agent_run", ReadOnly = false, Destructive = true, OpenWorld = false, UseStructuredContent = false)]
    [Description("Queue a natural-language request for asynchronous execution.")]
    public async Task<string> AgentRun(
        RequestContext<CallToolRequestParams> context,
        [Description("Non-empty natural-language work request.")] string request,
        [Description("Existing session UUID, or omit to create a session.")] Guid? session_id = null,
        [Description("Task-scoped temperature from 0.0 through 2.0.")] double? temperature = null,
        [Description("Optional PNG, JPEG, or WebP images.")] IReadOnlyList<McpImageInput>? images = null,
        CancellationToken cancellationToken = default)
    {
        RequireOnly(context, "request", "session_id", "temperature", "images");
        ValidateImages(images);
        Guid? mcpTaskId = GetRelatedTaskId(context.Params?.Meta) ?? GetRelatedTaskId(context.JsonRpcRequest) ??
            taskStore.ResolveTaskId(cancellationToken);
        TaskSnapshot task = registry.Enqueue(request, session_id, temperature, images?.Select(static image => image.DataBase64).ToArray(), mcpTaskId);
        string message = locale.Format("mcp.task_queued", new Dictionary<string, object?>
        {
            ["task"] = task.TaskId,
            ["session"] = task.SessionId
        });
        string queued = string.Join('\n',
            $"task_id: {task.TaskId}",
            $"session_id: {task.SessionId}",
            "status: queued",
            $"message: {message}");
        if (mcpTaskId is null) return queued;

        using CancellationTokenRegistration registration = cancellationToken.Register(() => registry.Cancel(task.TaskId));
        TaskSnapshot terminal = await registry.WaitForTerminalAsync(task.TaskId, cancellationToken).ConfigureAwait(false);
        if (terminal.State == AgentTaskState.Failed) throw new InvalidOperationException(terminal.Result ?? "Agent task failed.");
        if (terminal.State == AgentTaskState.Cancelled) throw new OperationCanceledException("Agent task was cancelled.", cancellationToken);
        return FormatStatus(terminal);
    }

    private void ValidateImages(IReadOnlyList<McpImageInput>? images)
    {
        if (images is null) return;
        if (images.Count > visionLimits.MaxCount) throw new ArgumentException($"Image count exceeds configured limit {visionLimits.MaxCount}.");
        long total = 0;
        foreach (McpImageInput image in images)
        {
            if (image.MimeType is not ("image/png" or "image/jpeg" or "image/webp")) throw new ArgumentException($"Unsupported image MIME type '{image.MimeType}'.");
            byte[] bytes;
            try { bytes = Convert.FromBase64String(image.DataBase64); }
            catch (FormatException ex) { throw new ArgumentException("Image data_base64 is invalid.", ex); }
            if (bytes.LongLength > visionLimits.MaxBytesEach) throw new ArgumentException($"Image exceeds configured per-image limit {visionLimits.MaxBytesEach} bytes.");
            total = checked(total + bytes.LongLength);
            if (total > visionLimits.MaxTotalBytes) throw new ArgumentException($"Images exceed configured total limit {visionLimits.MaxTotalBytes} bytes.");
        }
    }

    [McpServerTool(Name = "agent_status", ReadOnly = true, Destructive = false, OpenWorld = false, UseStructuredContent = false)]
    [Description("Get the current state and result of a queued agent task.")]
    public string AgentStatus(
        RequestContext<CallToolRequestParams> context,
        [Description("Task UUID returned by agent_run.")] Guid task_id,
        [Description("Optional session UUID used to verify task ownership.")] Guid? session_id = null)
    {
        RequireOnly(context, "task_id", "session_id");
        if (!registry.TryGet(task_id, session_id, out TaskSnapshot? task) || task is null)
        {
            return string.Join('\n',
                $"task_id: {task_id}",
                "status: failed",
                $"message: {locale.Format("mcp.task_not_found", new Dictionary<string, object?> { ["task"] = task_id })}");
        }

        return FormatStatus(task);
    }

    private string FormatStatus(TaskSnapshot task)
    {
        string state = task.State.ToString().ToLowerInvariant();
        bool terminal = task.State is AgentTaskState.Completed or AgentTaskState.Failed or AgentTaskState.Cancelled;
        return string.Join('\n',
            $"task_id: {task.TaskId}",
            $"session_id: {task.SessionId}",
            $"status: {state}",
            $"completed_tasks: {task.CompletedTasks}/{task.TotalTasks}",
            $"current_task: {task.CurrentTask ?? string.Empty}",
            $"token_usage: input={task.InputTokens} output={task.OutputTokens} total={task.InputTokens + task.OutputTokens}",
            $"critical_path_tokens: {task.CriticalPathTokens.ToString(System.Globalization.CultureInfo.InvariantCulture)}",
            $"message: {locale.Get("mcp.status." + state)}",
            terminal ? $"result: {task.Result}" : "result:");
    }

    private static Guid? GetRelatedTaskId(object? metadata)
    {
        if (metadata is null) return null;
        JsonElement json = JsonSerializer.SerializeToElement(metadata);
        return Find(json);

        static Guid? Find(JsonElement json)
        {
            if (json.ValueKind == JsonValueKind.Object)
            {
                foreach (JsonProperty property in json.EnumerateObject())
                {
                    if (property.NameEquals(TasksProtocol.MetaRelatedTask)) return Parse(property.Value);
                    Guid? nested = Find(property.Value);
                    if (nested is not null) return nested;
                }
            }
            else if (json.ValueKind == JsonValueKind.Array)
            {
                foreach (JsonElement item in json.EnumerateArray())
                {
                    Guid? nested = Find(item);
                    if (nested is not null) return nested;
                }
            }
            return null;
        }

        static Guid? Parse(JsonElement related)
        {
        if (related.ValueKind == JsonValueKind.String && Guid.TryParse(related.GetString(), out Guid direct)) return direct;
        if (related.ValueKind == JsonValueKind.Object && related.TryGetProperty("taskId", out JsonElement taskId) &&
            Guid.TryParse(taskId.GetString(), out Guid nested)) return nested;
        return null;
        }
    }

    private static void RequireOnly(RequestContext<CallToolRequestParams> context, params string[] allowed)
    {
        if (context.Params?.Arguments is null) return;
        HashSet<string> allowedSet = allowed.ToHashSet(StringComparer.Ordinal);
        string? unknown = context.Params.Arguments.Keys.FirstOrDefault(key => !allowedSet.Contains(key));
        if (unknown is not null) throw new ArgumentException($"Unknown tool property '{unknown}'.");
    }
}