// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

using Microsoft.Extensions.AI;
using ModelContextProtocol.Client;

namespace ExtAIBot;

// Owns the McpClient lifetime, lists tools at startup, and returns them wrapped
// with citation extraction so search results populate the CitationCollector.
internal sealed class McpToolSet : IAsyncDisposable
{
    private readonly McpClient _client;
    private readonly IList<McpClientTool> _tools;
    private readonly ILogger<McpToolSet> _logger;

    private McpToolSet(McpClient client, IList<McpClientTool> tools, ILogger<McpToolSet> logger)
    {
        _client = client;
        _tools = tools;
        _logger = logger;
    }

    public static async Task<McpToolSet> CreateAsync(ILogger<McpToolSet> logger, CancellationToken cancellationToken = default)
    {
        McpClient client = await McpClient.CreateAsync(
            new HttpClientTransport(new HttpClientTransportOptions
            {
                Endpoint = new Uri("https://learn.microsoft.com/api/mcp"),
                Name = "MSLearn",
                TransportMode = HttpTransportMode.StreamableHttp
            }),
            cancellationToken: cancellationToken);

        IList<McpClientTool> tools =
            await client.ListToolsAsync(cancellationToken: cancellationToken);

        return new McpToolSet(client, tools, logger);
    }

    // Returns each MCP tool wrapped so its results feed into the CitationCollector.
    public IList<AITool> GetTools(CitationCollector citations) =>
        [.. _tools.Select(t => new CitationCapturingTool(t, citations, _logger))];

    public ValueTask DisposeAsync() => _client.DisposeAsync();
}

internal sealed class McpToolSetLifetimeService(ILogger<McpToolSet> logger) : IHostedService
{
    private McpToolSet? _value;

    public McpToolSet Value => _value ?? throw new InvalidOperationException("MCP tool set is not initialized.");

    public async Task StartAsync(CancellationToken cancellationToken)
    {
        _value = await McpToolSet.CreateAsync(logger, cancellationToken);
    }

    public async Task StopAsync(CancellationToken cancellationToken)
    {
        if (_value is null) return;

        await _value.DisposeAsync();
        _value = null;
    }
}

// Wraps an McpClientTool, delegating all metadata to it while intercepting
// InvokeCoreAsync to extract citation data from the raw result string.
file sealed class CitationCapturingTool(McpClientTool inner, CitationCollector citations, ILogger logger)
    : DelegatingAIFunction(inner)
{
    protected override async ValueTask<object?> InvokeCoreAsync(
        AIFunctionArguments arguments,
        CancellationToken cancellationToken)
    {
        logger.LogInformation("[tool] {Name}({Args})",
            inner.Name,
            string.Join(", ", arguments.Select(a => $"{a.Key}={a.Value}")));

        object? result = await inner.InvokeAsync(arguments, cancellationToken);
        if (result?.ToString() is string text)
            citations.TryExtract(text);
        return result;
    }
}