using System.ComponentModel;
using Agent.Core.Localization;
using Agent.Core.Tasks;
using ModelContextProtocol.Server;
using ModelContextProtocol.Protocol;

namespace Agent.Mcp;

[McpServerToolType]
public sealed class AgentTools(TaskRegistry registry, LocaleCatalog locale, VisionLimits visionLimits)
{
    [McpServerTool(Name = "agent_run", ReadOnly = false, Destructive = true, OpenWorld = false, UseStructuredContent = false)]
    [Description("Queue a natural-language request for asynchronous execution.")]
    public 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)
    {
        RequireOnly(context, "request", "session_id", "temperature", "images");
        ValidateImages(images);
        TaskSnapshot task = registry.Enqueue(request, session_id, temperature, images?.Select(static image => image.DataBase64).ToArray());
        string message = locale.Format("mcp.task_queued", new Dictionary<string, object?>
        {
            ["task"] = task.TaskId,
            ["session"] = task.SessionId
        });
        return string.Join('\n',
            $"task_id: {task.TaskId}",
            $"session_id: {task.SessionId}",
            "status: queued",
            $"message: {message}");
    }

    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 })}");
        }

        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 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}'.");
    }
}