using System.Text.Json;
using System.Text.Json.Serialization;
using System.Text.Json.Serialization.Metadata;
using System.Text.Json.Nodes;
using System.Reflection;
using Agent.Core.Localization;
using Agent.Core.Sessions;
using Agent.Core.Tasks;
using Agent.Core.Ollama;
using Agent.Core.Runtime;
using Agent.Core.Tools;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using Microsoft.AspNetCore.Builder;
using Microsoft.AspNetCore.Hosting;
using Microsoft.AspNetCore.Http;
using ModelContextProtocol.Server;
using ModelContextProtocol.Extensions.Tasks;
using ModelContextProtocol;
using ModelContextProtocol.Protocol;
using System.Diagnostics;

namespace Agent.Mcp;

public static class McpHost
{
    public static async Task RunStdioAsync(LocaleCatalog locale, OllamaSettings ollamaSettings, AgentRuntimeSettings runtimeSettings, VisionLimits visionLimits, CancellationToken cancellationToken, InternalToolDispatcher? toolDispatcher = null, IAgentRuntimeObserver? observer = null, McpTaskStoreSettings? taskStoreSettings = null, IMcpHostObserver? hostObserver = null, IReadOnlyList<McpServerTool>? additionalTools = null, IReadOnlySet<string>? taskCapableTools = null, IReadOnlySet<string>? requiredTaskTools = null, Action<IMcpServerBuilder>? configureBuilder = null, SessionStore? sessionStore = null)
    {
        HostApplicationBuilder builder = Host.CreateApplicationBuilder(new HostApplicationBuilderSettings
        {
            Args = [],
            ContentRootPath = AppContext.BaseDirectory
        });
        builder.Logging.ClearProviders();
        var registry = new TaskRegistry(sessionStore);
        using var httpClient = new HttpClient { Timeout = TimeSpan.FromMinutes(10) };
        var ollama = new OllamaClient(httpClient, ollamaSettings);
        var runtime = new AgentRuntime(registry, ollama, runtimeSettings, toolDispatcher, observer, locale: locale);
        builder.Services.AddSingleton(locale);
        builder.Services.AddSingleton(registry);
        builder.Services.AddSingleton(runtime);
        builder.Services.AddHostedService<QueuedTaskWorker>();

        var taskStore = new PersistentMcpTaskStore(taskStoreSettings ?? DefaultTaskStoreSettings());
        McpServerTool[] tools = CreateTools(registry, locale, visionLimits, taskStore).Concat(additionalTools ?? []).ToArray();

        IMcpServerBuilder mcpBuilder = builder.Services
            .AddMcpServer()
            .WithStdioServerTransport()
            .WithTools(tools);
        ConfigureAgentTaskFilter(mcpBuilder);
        configureBuilder?.Invoke(mcpBuilder);
        mcpBuilder.WithTasks(taskStore, options => ConfigureTasks(options, taskCapableTools, requiredTaskTools));

        using IHost host = builder.Build();
        await host.StartAsync(cancellationToken).ConfigureAwait(false);
        hostObserver?.Ready(new McpReadyInfo("stdio", null, Environment.ProcessId, "2026-07-28", false, true));
        await host.WaitForShutdownAsync(cancellationToken).ConfigureAwait(false);
    }

    public static async Task RunHttpAsync(LocaleCatalog locale, OllamaSettings ollamaSettings, AgentRuntimeSettings runtimeSettings, VisionLimits visionLimits, string hostName, int port, CancellationToken cancellationToken, InternalToolDispatcher? toolDispatcher = null, IAgentRuntimeObserver? observer = null, McpTaskStoreSettings? taskStoreSettings = null, IMcpHostObserver? hostObserver = null, IReadOnlyList<McpServerTool>? additionalTools = null, IReadOnlySet<string>? taskCapableTools = null, IReadOnlySet<string>? requiredTaskTools = null, Action<IMcpServerBuilder>? configureBuilder = null, SessionStore? sessionStore = null)
    {
        if (!System.Net.IPAddress.TryParse(hostName, out System.Net.IPAddress? address) || !System.Net.IPAddress.IsLoopback(address))
            throw new ArgumentException("HTTP transport host must be a loopback IP literal.", nameof(hostName));
        WebApplicationBuilder builder = WebApplication.CreateBuilder(new WebApplicationOptions { Args = [], ContentRootPath = AppContext.BaseDirectory });
        builder.Logging.ClearProviders();
        string listenHost = address.AddressFamily == System.Net.Sockets.AddressFamily.InterNetworkV6 ? $"[{address}]" : address.ToString();
        builder.WebHost.UseUrls($"http://{listenHost}:{port}");
        var registry = new TaskRegistry(sessionStore);
        var httpClient = new HttpClient { Timeout = TimeSpan.FromMinutes(10) };
        var runtime = new AgentRuntime(registry, new OllamaClient(httpClient, ollamaSettings), runtimeSettings, toolDispatcher, observer, locale: locale);
        builder.Services.AddSingleton(locale);
        builder.Services.AddSingleton(registry);
        builder.Services.AddSingleton(runtime);
        builder.Services.AddHostedService<QueuedTaskWorker>();
        var taskStore = new PersistentMcpTaskStore(taskStoreSettings ?? DefaultTaskStoreSettings());
        McpServerTool[] tools = CreateTools(registry, locale, visionLimits, taskStore).Concat(additionalTools ?? []).ToArray();
        IMcpServerBuilder mcpBuilder = builder.Services.AddMcpServer()
            .WithHttpTransport(options => options.Stateless = true)
            .WithTools(tools);
        ConfigureAgentTaskFilter(mcpBuilder);
        configureBuilder?.Invoke(mcpBuilder);
        mcpBuilder.WithTasks(taskStore, options => ConfigureTasks(options, taskCapableTools, requiredTaskTools));
        await using WebApplication app = builder.Build();
        app.Use(async (context, next) =>
        {
            var stopwatch = Stopwatch.StartNew();
            DateTimeOffset timestamp = DateTimeOffset.UtcNow;
            hostObserver?.HttpRequestStarted(new McpHttpRequestStartInfo(context.Connection.Id, context.TraceIdentifier,
                context.Connection.RemoteIpAddress?.ToString(), context.Connection.RemotePort,
                context.Request.Method, context.Request.Path.Value ?? string.Empty));
            try
            {
                if (context.Request.Headers.TryGetValue("Origin", out Microsoft.Extensions.Primitives.StringValues origins) &&
                    origins.Any(origin => !IsLoopbackOrigin(origin)))
                {
                    context.Response.StatusCode = StatusCodes.Status403Forbidden;
                }
                else await next(context).ConfigureAwait(false);
            }
            finally
            {
                stopwatch.Stop();
                hostObserver?.HttpRequest(new McpHttpRequestInfo(timestamp, context.Connection.Id, context.TraceIdentifier,
                    context.Connection.RemoteIpAddress?.ToString(), context.Connection.RemotePort,
                    context.Connection.LocalIpAddress?.ToString(), context.Connection.LocalPort,
                    context.Request.Method, context.Request.Path.Value ?? string.Empty, context.Response.StatusCode,
                    stopwatch.Elapsed.TotalMilliseconds, context.Request.ContentType, context.Request.ContentLength,
                    context.Request.Headers.UserAgent.ToString(), context.Request.Headers["MCP-Protocol-Version"].ToString(),
                    context.Request.Headers["MCP-Method"].ToString(), context.Request.Headers["MCP-Name"].ToString()));
            }
        });
        app.MapMcp("/mcp");
        try
        {
            await app.StartAsync(cancellationToken).ConfigureAwait(false);
            hostObserver?.Ready(new McpReadyInfo("http", $"http://{listenHost}:{port}/mcp", Environment.ProcessId, "2026-07-28", true, true));
            await app.WaitForShutdownAsync(cancellationToken).ConfigureAwait(false);
        }
        finally { httpClient.Dispose(); }
    }

    private static McpTaskStoreSettings DefaultTaskStoreSettings() =>
        new(Path.Combine(Path.GetTempPath(), "agent-mcp", "mcp-tasks"), TimeSpan.FromHours(24), 1000);

    private static void ConfigureTasks(McpTasksOptions options, IReadOnlySet<string>? taskCapableTools, IReadOnlySet<string>? requiredTaskTools)
    {
        IReadOnlySet<string> names = taskCapableTools ?? new HashSet<string>(["agent_run"], StringComparer.Ordinal);
        options.ExecutionModeSelector = context =>
            context.MatchedPrimitive is not McpServerTool tool ? McpTaskExecutionMode.Synchronous :
            requiredTaskTools?.Contains(tool.ProtocolTool.Name) == true ? McpTaskExecutionMode.Required :
            names.Contains(tool.ProtocolTool.Name) ? McpTaskExecutionMode.Optional : McpTaskExecutionMode.Synchronous;
    }

    private static void ConfigureAgentTaskFilter(IMcpServerBuilder builder)
    {
        builder.WithRequestFilters(filters => filters.AddCallToolFilter(next => async (context, cancellationToken) =>
        {
            CallToolResult result = await next(context, cancellationToken).ConfigureAwait(false);
            if (context.Params?.Name == "agent_run" && result.IsError == true)
                throw new McpProtocolException("The Agent task failed.", McpErrorCode.InternalError);
            return result;
        }));
    }

    private static bool IsLoopbackOrigin(string? origin)
    {
        return Uri.TryCreate(origin, UriKind.Absolute, out Uri? uri) &&
            (uri.Scheme == Uri.UriSchemeHttp || uri.Scheme == Uri.UriSchemeHttps) &&
            (string.Equals(uri.Host, "localhost", StringComparison.OrdinalIgnoreCase) ||
             (System.Net.IPAddress.TryParse(uri.Host, out System.Net.IPAddress? address) && System.Net.IPAddress.IsLoopback(address)));
    }

    private static McpServerTool[] CreateTools(TaskRegistry registry, LocaleCatalog locale, VisionLimits visionLimits, PersistentMcpTaskStore taskStore)
    {
        var serializerOptions = new JsonSerializerOptions(JsonSerializerDefaults.Web)
        {
            PropertyNamingPolicy = JsonNamingPolicy.SnakeCaseLower,
            TypeInfoResolver = new DefaultJsonTypeInfoResolver(),
            UnmappedMemberHandling = JsonUnmappedMemberHandling.Disallow
        };
        var target = new AgentTools(registry, locale, visionLimits, taskStore);
        McpServerTool[] tools = typeof(AgentTools)
            .GetMethods(BindingFlags.Instance | BindingFlags.Public | BindingFlags.DeclaredOnly)
            .Where(static method => method.GetCustomAttribute<McpServerToolAttribute>() is not null)
            .Select(method => McpServerTool.Create(method, target, new McpServerToolCreateOptions { SerializerOptions = serializerOptions }))
            .ToArray();
        foreach (McpServerTool tool in tools) MakeSchemaStrict(tool);
        return tools;
    }

    private static void MakeSchemaStrict(McpServerTool tool)
    {
        JsonNode root = JsonNode.Parse(tool.ProtocolTool.InputSchema.GetRawText())
            ?? throw new InvalidOperationException($"Tool '{tool.ProtocolTool.Name}' has no input schema.");
        Visit(root);
        tool.ProtocolTool.InputSchema = JsonSerializer.SerializeToElement(root);

        static void Visit(JsonNode? node)
        {
            if (node is JsonObject obj)
            {
                if (obj["type"] is JsonValue type && type.TryGetValue(out string? typeName) && typeName == "object")
                    obj["additionalProperties"] = false;
                foreach (KeyValuePair<string, JsonNode?> property in obj.ToArray()) Visit(property.Value);
            }
            else if (node is JsonArray array)
            {
                foreach (JsonNode? item in array) Visit(item);
            }
        }
    }
}