using System.Diagnostics;
using System.Text.Json;
using Microsoft.Extensions.AI;
using Svrnty.CQRS.Abstractions;
using Svrnty.Sample.AI.Tools;
using Svrnty.Sample.Data;
using Svrnty.Sample.Data.Entities;

namespace Svrnty.Sample.AI.Commands;

/// <summary>
/// Handler for executing AI agent commands with function calling support and full observability
/// </summary>
public class ExecuteAgentCommandHandler(
    IChatClient chatClient,
    AgentDbContext dbContext,
    MathTool mathTool,
    DatabaseQueryTool dbTool,
    ILogger<ExecuteAgentCommandHandler> logger) : ICommandHandler<ExecuteAgentCommand, AgentResponse>
{
    private static readonly ActivitySource ActivitySource = new("Svrnty.AI.Agent");
    private const int MaxFunctionCallIterations = 10; // Prevent infinite loops

    public async Task<AgentResponse> HandleAsync(
        ExecuteAgentCommand command,
        CancellationToken cancellationToken = default)
    {
        var conversationId = Guid.NewGuid();

        // Start root trace
        using var activity = ActivitySource.StartActivity("agent.execute", ActivityKind.Server);
        activity?.SetTag("agent.conversation_id", conversationId);
        activity?.SetTag("agent.prompt", command.Prompt);
        activity?.SetTag("agent.model", "qwen2.5-coder:7b");

        try
        {
            var messages = new List<ChatMessage>
            {
                new(ChatRole.User, command.Prompt)
            };

            // Register available tools
            using (var toolActivity = ActivitySource.StartActivity("tools.register"))
            {
                var tools = new List<AIFunction>
                {
                    AIFunctionFactory.Create(mathTool.Add),
                    AIFunctionFactory.Create(mathTool.Multiply),
                    AIFunctionFactory.Create(dbTool.GetMonthlyRevenue),
                    AIFunctionFactory.Create(dbTool.GetRevenueRange),
                    AIFunctionFactory.Create(dbTool.CountCustomersByState),
                    AIFunctionFactory.Create(dbTool.CountCustomersByTier),
                    AIFunctionFactory.Create(dbTool.GetCustomers)
                };

                toolActivity?.SetTag("tools.count", tools.Count);
                toolActivity?.SetTag("tools.names", string.Join(",", tools.Select(t => t.Metadata.Name)));

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

                var functionLookup = tools.ToDictionary(
                    f => f.Metadata.Name,
                    f => f,
                    StringComparer.OrdinalIgnoreCase
                );

                // Initial AI completion
                using (var llmActivity = ActivitySource.StartActivity("llm.completion"))
                {
                    llmActivity?.SetTag("llm.iteration", 0);
                    var completion = await chatClient.CompleteAsync(messages, options, cancellationToken);
                    messages.Add(completion.Message);

                    // Function calling loop
                    var iterations = 0;
                    while (completion.Message.Contents.OfType<FunctionCallContent>().Any()
                           && iterations < MaxFunctionCallIterations)
                    {
                        iterations++;

                        foreach (var functionCall in completion.Message.Contents.OfType<FunctionCallContent>())
                        {
                            using var funcActivity = ActivitySource.StartActivity($"function.{functionCall.Name}");
                            funcActivity?.SetTag("function.name", functionCall.Name);
                            funcActivity?.SetTag("function.arguments", JsonSerializer.Serialize(functionCall.Arguments));

                            try
                            {
                                if (!functionLookup.TryGetValue(functionCall.Name, out var function))
                                {
                                    throw new InvalidOperationException($"Function '{functionCall.Name}' not found");
                                }

                                var result = await function.InvokeAsync(functionCall.Arguments, cancellationToken);
                                funcActivity?.SetTag("function.result", result?.ToString() ?? "null");
                                funcActivity?.SetTag("function.success", true);

                                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)
                            {
                                funcActivity?.SetTag("function.success", false);
                                funcActivity?.SetTag("error.message", ex.Message);

                                var errorMessage = new ChatMessage(ChatRole.Tool, $"Error: {ex.Message}");
                                errorMessage.Contents.Add(new FunctionResultContent(functionCall.CallId, functionCall.Name, $"Error: {ex.Message}"));
                                messages.Add(errorMessage);
                            }
                        }

                        using (var nextLlmActivity = ActivitySource.StartActivity("llm.completion"))
                        {
                            nextLlmActivity?.SetTag("llm.iteration", iterations);
                            completion = await chatClient.CompleteAsync(messages, options, cancellationToken);
                            messages.Add(completion.Message);
                        }
                    }

                    // Store conversation in database
                    var conversation = new Conversation
                    {
                        Id = conversationId,
                        Messages = messages.Select(m => new ConversationMessage
                        {
                            Role = m.Role.ToString(),
                            Content = m.Text ?? string.Empty,
                            Timestamp = DateTime.UtcNow
                        }).ToList()
                    };

                    dbContext.Conversations.Add(conversation);
                    await dbContext.SaveChangesAsync(cancellationToken);

                    activity?.SetTag("agent.success", true);
                    activity?.SetTag("agent.iterations", iterations);
                    activity?.SetTag("agent.response_preview", completion.Message.Text?.Substring(0, Math.Min(100, completion.Message.Text.Length)));

                    logger.LogInformation("Agent executed successfully for conversation {ConversationId}", conversationId);

                    return new AgentResponse(
                        Content: completion.Message.Text ?? "No response",
                        ConversationId: conversationId
                    );
                }
            }
        }
        catch (Exception ex)
        {
            activity?.SetTag("agent.success", false);
            activity?.SetTag("error.type", ex.GetType().Name);
            activity?.SetTag("error.message", ex.Message);

            logger.LogError(ex, "Agent execution failed for conversation {ConversationId}", conversationId);
            throw;
        }
    }
}