using System.Text.Json;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Logging;
using TeamUp.Modules.Integrations.Persistence;
using TeamUp.Modules.Integrations.Security;
using TeamUp.SharedKernel.Ai;

namespace TeamUp.Modules.Integrations.Mcp;

/// <summary>
/// Resolves an org's MCP server configs, decrypts their headers server-side, and talks to them via
/// <see cref="McpClient"/>. Discovery is best-effort: a server that fails to connect is logged and
/// skipped so it never fails the agent run. The decrypted headers never leave the server.
/// </summary>
internal sealed class McpGateway(
    IntegrationsDbContext db,
    McpClient client,
    ISecretProtector protector,
    ILogger<McpGateway> logger) : IMcpGateway
{
    public async Task<IReadOnlyList<McpToolDescriptor>> ListToolsAsync(
        Guid organizationId, IReadOnlyCollection<Guid> serverIds, CancellationToken cancellationToken = default)
    {
        if (serverIds.Count == 0)
        {
            return [];
        }

        var idSet = serverIds.ToHashSet();
        var servers = await db.McpServers
            .Where(s => s.OrganizationId == organizationId && s.Enabled && idSet.Contains(s.Id))
            .ToListAsync(cancellationToken);

        var tools = new List<McpToolDescriptor>();
        foreach (var server in servers)
        {
            try
            {
                var discovered = await client.ListToolsAsync(server.Endpoint, DecryptHeaders(server.EncryptedHeaders), cancellationToken);
                tools.AddRange(discovered.Select(t =>
                    new McpToolDescriptor(server.Id, server.Name, t.Name, t.Description, t.InputSchemaJson)));
            }
            catch (Exception ex)
            {
                logger.LogWarning(ex, "MCP server {Server} ({Endpoint}) unreachable; skipping its tools.", server.Name, server.Endpoint);
            }
        }

        return tools;
    }

    public async Task<McpToolResult> CallToolAsync(
        Guid organizationId, Guid serverId, string toolName, string argumentsJson, CancellationToken cancellationToken = default)
    {
        var server = await db.McpServers.FirstOrDefaultAsync(
            s => s.Id == serverId && s.OrganizationId == organizationId && s.Enabled, cancellationToken);
        if (server is null)
        {
            return new McpToolResult(false, null, "MCP server not found or disabled.");
        }

        try
        {
            var (success, content, error) = await client.CallToolAsync(
                server.Endpoint, DecryptHeaders(server.EncryptedHeaders), toolName, argumentsJson, cancellationToken);
            return new McpToolResult(success, content, error);
        }
        catch (Exception ex)
        {
            return new McpToolResult(false, null, ex.Message);
        }
    }

    private Dictionary<string, string>? DecryptHeaders(string? encrypted) =>
        string.IsNullOrEmpty(encrypted)
            ? null
            : JsonSerializer.Deserialize<Dictionary<string, string>>(protector.Unprotect(encrypted));
}