using OpenHarbor.MCP.Client.Core.Abstractions;
using OpenHarbor.MCP.Client.Core.Exceptions;
using OpenHarbor.MCP.Client.Core.Models;

namespace OpenHarbor.MCP.Client.Infrastructure;

/// <summary>
/// Implements MCP client for managing multiple MCP server connections.
/// Provides tool discovery and execution across configured servers.
/// </summary>
public class McpClient : IMcpClient
{
    private readonly List<McpServerConfig> _configs;
    private readonly Dictionary<string, IMcpServerConnection> _connections;

    public McpClient(List<McpServerConfig> configs)
    {
        _configs = configs ?? throw new ArgumentNullException(nameof(configs));
        _connections = new Dictionary<string, IMcpServerConnection>();
    }

    public async Task ConnectAllAsync(CancellationToken cancellationToken = default)
    {
        foreach (var config in _configs.Where(c => c.Enabled))
        {
            if (_connections.ContainsKey(config.Name))
            {
                continue;
            }

            var connection = CreateConnection(config);
            await connection.ConnectAsync(cancellationToken);
            _connections[config.Name] = connection;
        }
    }

    public async Task DisconnectAllAsync(CancellationToken cancellationToken = default)
    {
        foreach (var connection in _connections.Values)
        {
            await connection.DisconnectAsync(cancellationToken);
        }

        _connections.Clear();
    }

    public Task<IEnumerable<McpServerConfig>> GetConnectedServersAsync(
        CancellationToken cancellationToken = default)
    {
        var connectedConfigs = _configs
            .Where(c => _connections.ContainsKey(c.Name))
            .ToList();

        return Task.FromResult<IEnumerable<McpServerConfig>>(connectedConfigs);
    }

    public async Task<IEnumerable<McpTool>> ListToolsAsync(
        string serverName,
        CancellationToken cancellationToken = default)
    {
        var connection = GetConnection(serverName);
        return await connection.ListToolsAsync(cancellationToken);
    }

    public async Task<McpToolResult> CallToolAsync(
        string serverName,
        string toolName,
        Dictionary<string, object>? arguments = null,
        CancellationToken cancellationToken = default)
    {
        var connection = GetConnection(serverName);
        return await connection.CallToolAsync(toolName, arguments, cancellationToken);
    }

    public async Task PingAsync(string serverName, CancellationToken cancellationToken = default)
    {
        var connection = GetConnection(serverName);
        await connection.PingAsync(cancellationToken);
    }

    public async ValueTask DisposeAsync()
    {
        await DisconnectAllAsync();

        foreach (var connection in _connections.Values)
        {
            await connection.DisposeAsync();
        }

        _connections.Clear();
    }

    private IMcpServerConnection GetConnection(string serverName)
    {
        if (!_connections.TryGetValue(serverName, out var connection))
        {
            throw new ServerNotFoundException(serverName);
        }

        return connection;
    }

    private static IMcpServerConnection CreateConnection(McpServerConfig config)
    {
        return config.Transport switch
        {
            StdioTransportConfig => new StdioServerConnection(config),
            _ => throw new NotSupportedException(
                $"Transport type '{config.Transport.Type}' is not supported")
        };
    }
}