using Microsoft.Extensions.AI;

namespace Kds.Ai.AnomalyDetector.Mcp;

/// <summary>
/// Provides MCP tools (e.g., mysql_query) as AITool objects for agent use.
/// </summary>
public interface IMcpToolProvider
{
    /// <summary>
    /// Retrieves available tools as AITool objects
    /// that can be passed directly to an agent.
    /// </summary>
    Task<IList<AITool>> GetToolsAsync(CancellationToken ct = default);
}

using Kds.Ai.AnomalyDetector.Settings;
using Microsoft.Extensions.AI;
using ModelContextProtocol.Client;

namespace Kds.Ai.AnomalyDetector.Mcp;

/// <summary>
/// Manages the MCP client lifecycle and provides tools from the MySQL MCP server.
/// </summary>
internal sealed class McpToolProvider(McpMySqlSettings settings) : IMcpToolProvider, IAsyncDisposable
{
    private McpClient? _client;

    /// <summary>
    /// Creates the MCP client and connects to the MySQL MCP server.
    /// </summary>
    public async Task<McpClient> GetClientAsync(CancellationToken ct = default)
    {
        if (_client is not null)
        {
            return _client;
        }

        var transport = new StdioClientTransport(new StdioClientTransportOptions
        {
            Name = "mysql-mcp-server",
            Command = settings.Command,
            Arguments = settings.Arguments,
            EnvironmentVariables = settings.EnvironmentVariables
        });

        _client = await McpClient.CreateAsync(transport, cancellationToken: ct);
        return _client;
    }

    /// <inheritdoc />
    public async Task<IList<AITool>> GetToolsAsync(CancellationToken ct = default)
    {
        var client = await GetClientAsync(ct);
        var tools = await client.ListToolsAsync(cancellationToken: ct);
        return [.. tools.Cast<AITool>()];
    }

    public async ValueTask DisposeAsync()
    {
        if (_client is not null)
        {
            await _client.DisposeAsync();
            _client = null;
        }
    }
}

// Change the constructor parameter type:
public sealed class AnomalyDetectorAgent(
    IChatClient chatClient,
    IMcpToolProvider mcpToolProvider,  // was: McpToolProvider
    ILogger<AnomalyDetectorAgent> logger) : IAnomalyDetectorAgent

builder.Services.AddSingleton<IMcpToolProvider>(provider =>
{
    var settings = provider.GetRequiredService<IOptions<McpMySqlSettings>>();
    return new McpToolProvider(settings.Value);
});

using Kds.Ai.AnomalyDetector.Mcp;
using Microsoft.Extensions.AI;

namespace Kds.Ai.AnomalyDetector.Tests.Fakes;

/// <summary>
/// Fake MCP tool provider that returns canned SQL query results
/// without needing a real MCP server process.
/// </summary>
public sealed class FakeMcpToolProvider : IMcpToolProvider
{
    private readonly Dictionary<string, Func<string, string>> _queryHandlers = new();
    
    /// <summary>
    /// Register a handler for SQL queries matching a pattern.
    /// The handler receives the full SQL query and returns a JSON result.
    /// </summary>
    public FakeMcpToolProvider OnQuery(string pattern, Func<string, string> handler)
    {
        _queryHandlers[pattern] = handler;
        return this;
    }
    
    /// <summary>
    /// Register a handler that returns a fixed result for queries matching a pattern.
    /// </summary>
    public FakeMcpToolProvider OnQuery(string pattern, string fixedResult)
    {
        _queryHandlers[pattern] = _ => fixedResult;
        return this;
    }
    
    /// <inheritdoc />
    public Task<IList<AITool>> GetToolsAsync(CancellationToken ct = default)
    {
        var tool = AIFunctionFactory.Create(
            (string query) =>
            {
                foreach (var (pattern, handler) in _queryHandlers)
                {
                    if (query.Contains(pattern, StringComparison.OrdinalIgnoreCase))
                    {
                        return handler(query);
                    }
                }
                
                // Default: return empty result set
                return "[]";
            },
            name: "mysql_query",
            description: "Execute a SQL query against the MySQL database");
        
        return Task.FromResult<IList<AITool>>([tool]);
    }
}

using Microsoft.Extensions.AI;

namespace Kds.Ai.AnomalyDetector.Tests.Fakes;

/// <summary>
/// Fake IChatClient that returns scripted responses without calling an LLM.
/// Supports function calling by detecting tool call requests and routing them.
/// </summary>
public sealed class FakeChatClient : IChatClient
{
    private readonly Queue<ChatResponse> _responses = new();
    
    /// <summary>
    /// Enqueue a text response (the agent will receive this as assistant message)
    /// </summary>
    public FakeChatClient EnqueueTextResponse(string text)
    {
        _responses.Enqueue(new ChatResponse(new ChatMessage(ChatRole.Assistant, text)));
        return this;
    }
    
    /// <summary>
    /// Enqueue a function call response (simulates agent requesting a tool call)
    /// </summary>
    public FakeChatClient EnqueueFunctionCall(string functionName, string argumentsJson)
    {
        var functionCall = new FunctionCallContent(
            callId: Guid.NewGuid().ToString(),
            name: functionName,
            arguments: System.Text.Json.JsonSerializer.Deserialize<Dictionary<string, object?>>(argumentsJson));
        
        var message = new ChatMessage(ChatRole.Assistant, [functionCall]);
        _responses.Enqueue(new ChatResponse(message));
        return this;
    }
    
    public Task<ChatResponse> GetResponseAsync(
        IEnumerable<ChatMessage> messages,
        ChatOptions? options = null,
        CancellationToken cancellationToken = default)
    {
        if (_responses.Count == 0)
        {
            throw new InvalidOperationException(
                "FakeChatClient has no more scripted responses. " +
                "Enqueue more responses or check the test setup.");
        }
        
        return Task.FromResult(_responses.Dequeue());
    }
    
    public IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
        IEnumerable<ChatMessage> messages,
        ChatOptions? options = null,
        CancellationToken cancellationToken = default)
    {
        throw new NotSupportedException("Streaming not supported in fake");
    }
    
    public ChatClientMetadata Metadata => new("FakeChatClient");
    
    public void Dispose() { }
}

using Kds.Ai.Data.Enums;

namespace Kds.Ai.AnomalyDetector.Tests.Fakes;

/// <summary>
/// Pre-built scenarios for testing.
/// Each scenario configures a FakeMcpToolProvider and FakeChatClient
/// to simulate a specific detection outcome.
/// </summary>
public static class ScriptedAgentScenarios
{
    /// <summary>
    /// Scenario: Agent queries data, finds no anomalies, calls emit_conclusion with anomalyDetected=false
    /// </summary>
    public static (FakeMcpToolProvider Tools, FakeChatClient Chat) NoAnomaly()
    {
        var tools = new FakeMcpToolProvider()
            .OnQuery("SELECT", "[]"); // All queries return empty
        
        var chat = new FakeChatClient()
            // First: agent requests mysql_query
            .EnqueueFunctionCall("mysql_query", """{"query": "SELECT COUNT(*) as cnt FROM orders WHERE done=1"}""")
            // Second: after getting tool result, agent calls emit_conclusion
            .EnqueueFunctionCall("emit_conclusion", """
            {
                "anomalyDetected": false,
                "summary": "No anomalies detected in the analyzed period",
                "findings": []
            }
            """);
        
        return (tools, chat);
    }
    
    /// <summary>
    /// Scenario: Agent queries data, finds anomaly, calls emit_conclusion with findings
    /// </summary>
    public static (FakeMcpToolProvider Tools, FakeChatClient Chat) SingleAnomaly(
        EventSeverity severity = EventSeverity.Medium)
    {
        var tools = new FakeMcpToolProvider()
            .OnQuery("SELECT", """[{"hour": 12, "avg_prep": 450, "baseline": 280}]""");
        
        var chat = new FakeChatClient()
            .EnqueueFunctionCall("mysql_query", """{"query": "SELECT AVG(preparation_time) FROM orders"}""")
            .EnqueueFunctionCall("emit_conclusion", $$"""
            {
                "anomalyDetected": true,
                "severity": "{{severity}}",
                "summary": "Significant deviation detected during lunch hours",
                "detailedAnalysis": "Analysis found 60% increase in average values during 12:00-13:00 compared to 28-day baseline.",
                "findings": [
                    {
                        "description": "60% increase during 12:00-13:00",
                        "severity": "{{severity}}",
                        "affectedPeriod": "2025-01-15 12:00-13:00",
                        "evidence": "Average 450s vs baseline 280s (2.8 standard deviations)"
                    }
                ]
            }
            """);
        
        return (tools, chat);
    }
    
    /// <summary>
    /// Scenario: Agent encounters an error during data access
    /// </summary>
    public static (FakeMcpToolProvider Tools, FakeChatClient Chat) DataError()
    {
        var tools = new FakeMcpToolProvider()
            .OnQuery("SELECT", _ => throw new InvalidOperationException("Connection refused"));
        
        var chat = new FakeChatClient()
            .EnqueueFunctionCall("mysql_query", """{"query": "SELECT 1"}""");
        
        return (tools, chat);
    }
}

using Kds.Ai.AnomalyDetector.Mcp;
using Microsoft.Extensions.AI;
using Microsoft.Extensions.DependencyInjection;

namespace Kds.Ai.AnomalyDetector.Tests.Fakes;

public static class TestingExtensions
{
    /// <summary>
    /// Replace real services with fakes for testing.
    /// Uses a scripted scenario to configure the fake chat client and MCP tools.
    /// </summary>
    public static IServiceCollection UseFakes(
        this IServiceCollection services,
        FakeChatClient fakeChatClient,
        FakeMcpToolProvider fakeMcpToolProvider)
    {
        // Replace IChatClient with fake
        services.AddSingleton<IChatClient>(fakeChatClient);
        
        // Replace IMcpToolProvider with fake
        services.AddSingleton<IMcpToolProvider>(fakeMcpToolProvider);
        
        return services;
    }
}

[Fact]
public async Task DetectsAnomaly_WhenDataShowsDeviation()
{
    var (tools, chat) = ScriptedAgentScenarios.SingleAnomaly(EventSeverity.High);
    
    // Create agent with fakes
    var agent = new AnomalyDetectorAgent(chat, tools, NullLogger<AnomalyDetectorAgent>.Instance);
    
    var spec = new MonitorSpec
    {
        Name = "Test Monitor",
        SystemPrompt = "Analyze data and detect anomalies",
        State = MonitorState.Active
    };
    
    var result = await agent.RunDetectionAsync(spec, DateTimeOffset.UtcNow);
    
    Assert.True(result.AnomalyDetected);
    Assert.Single(result.Findings);
    Assert.Contains("12:00", result.Findings[0].AffectedPeriod);
}

// In Program.cs, add a --mock flag
if (args.Contains("--mock"))
{
    var (tools, chat) = ScriptedAgentScenarios.SingleAnomaly();
    builder.Services.UseFakes(chat, tools);
}