using Microsoft.Extensions.AI;
using Svrnty.CQRS.Abstractions;
using Svrnty.Sample.AI.Tools;

namespace Svrnty.Sample.AI.Commands;

/// <summary>
/// Handler for executing AI agent commands with function calling support
/// </summary>
public class ExecuteAgentCommandHandler(IChatClient chatClient) : ICommandHandler<ExecuteAgentCommand, AgentResponse>
{
    // In-memory conversation store (replace with proper persistence in production)
    private static readonly Dictionary<Guid, List<ChatMessage>> ConversationStore = new();
    private const int MaxFunctionCallIterations = 10; // Prevent infinite loops

    public async Task<AgentResponse> HandleAsync(
        ExecuteAgentCommand command,
        CancellationToken cancellationToken = default)
    {
        var conversationId = Guid.NewGuid();
        var messages = new List<ChatMessage>
        {
            new(ChatRole.User, command.Prompt)
        };

        // Register available tools
        var mathTool = new MathTool();
        var dbTool = new DatabaseQueryTool();
        var tools = new List<AIFunction>
        {
            // Math tools
            AIFunctionFactory.Create(mathTool.Add),
            AIFunctionFactory.Create(mathTool.Multiply),

            // Business tools
            AIFunctionFactory.Create(dbTool.GetMonthlyRevenue),
            AIFunctionFactory.Create(dbTool.GetRevenueRange),
            AIFunctionFactory.Create(dbTool.CountCustomersByState),
            AIFunctionFactory.Create(dbTool.CountCustomersByTier),
            AIFunctionFactory.Create(dbTool.GetCustomers)
        };

        var options = new ChatOptions
        {
            ModelId = "qwen2.5-coder:7b",
            Tools = tools.Cast<AITool>().ToList()
        };

        // Create function lookup by name for invocation
        var functionLookup = tools.ToDictionary(
            f => f.Metadata.Name,
            f => f,
            StringComparer.OrdinalIgnoreCase
        );

        // Initial AI completion
        var completion = await chatClient.CompleteAsync(messages, options, cancellationToken);
        messages.Add(completion.Message);

        // Function calling loop - continue until no more function calls or max iterations
        var iterations = 0;
        while (completion.Message.Contents.OfType<FunctionCallContent>().Any() && iterations < MaxFunctionCallIterations)
        {
            iterations++;

            // Execute all function calls from the response
            foreach (var functionCall in completion.Message.Contents.OfType<FunctionCallContent>())
            {
                try
                {
                    // Look up the actual function
                    if (!functionLookup.TryGetValue(functionCall.Name, out var function))
                    {
                        throw new InvalidOperationException($"Function '{functionCall.Name}' not found");
                    }

                    // Invoke the function with arguments
                    var result = await function.InvokeAsync(functionCall.Arguments, cancellationToken);

                    // Add function result to conversation as a tool message
                    var toolMessage = new ChatMessage(ChatRole.Tool, result?.ToString() ?? "null");
                    toolMessage.Contents.Add(new FunctionResultContent(functionCall.CallId, functionCall.Name, result));
                    messages.Add(toolMessage);
                }
                catch (Exception ex)
                {
                    // Handle function call errors gracefully
                    var errorMessage = new ChatMessage(ChatRole.Tool, $"Error executing {functionCall.Name}: {ex.Message}");
                    errorMessage.Contents.Add(new FunctionResultContent(functionCall.CallId, functionCall.Name, $"Error: {ex.Message}"));
                    messages.Add(errorMessage);
                }
            }

            // Get next completion with function results
            completion = await chatClient.CompleteAsync(messages, options, cancellationToken);
            messages.Add(completion.Message);
        }

        // Store conversation for potential future use
        ConversationStore[conversationId] = messages;

        return new AgentResponse(
            Content: completion.Message.Text ?? "No response",
            ConversationId: conversationId
        );
    }
}