using System;
using System.Threading;
using System.Threading.Tasks;
using Codex.Dal;
using Codex.Dal.Enums;
using FluentValidation;
using Microsoft.EntityFrameworkCore;
using OpenHarbor.CQRS.Abstractions;

namespace Codex.CQRS.Commands;

/// <summary>
/// Completes an agent execution with results and metrics
/// </summary>
public record CompleteAgentExecutionCommand
{
    /// <summary>Execution ID to complete</summary>
    public Guid ExecutionId { get; init; }

    /// <summary>Agent's output/response</summary>
    public string Output { get; init; } = string.Empty;

    /// <summary>Execution status (Completed, Failed, Cancelled)</summary>
    public ExecutionStatus Status { get; init; }

    /// <summary>Input tokens consumed</summary>
    public int? InputTokens { get; init; }

    /// <summary>Output tokens generated</summary>
    public int? OutputTokens { get; init; }

    /// <summary>Estimated cost in USD</summary>
    public decimal? EstimatedCost { get; init; }

    /// <summary>Tool calls made (JSON array)</summary>
    public string? ToolCalls { get; init; }

    /// <summary>Tool call results (JSON array)</summary>
    public string? ToolCallResults { get; init; }

    /// <summary>Error message if failed</summary>
    public string? ErrorMessage { get; init; }
}

public class CompleteAgentExecutionCommandHandler : ICommandHandler<CompleteAgentExecutionCommand>
{
    private readonly CodexDbContext _dbContext;

    public CompleteAgentExecutionCommandHandler(CodexDbContext dbContext)
    {
        _dbContext = dbContext;
    }

    public async Task HandleAsync(CompleteAgentExecutionCommand command, CancellationToken cancellationToken = default)
    {
        var execution = await _dbContext.AgentExecutions
            .FirstOrDefaultAsync(e => e.Id == command.ExecutionId, cancellationToken);

        if (execution == null)
        {
            throw new InvalidOperationException($"Execution with ID {command.ExecutionId} not found");
        }

        if (execution.Status != ExecutionStatus.Running)
        {
            throw new InvalidOperationException($"Execution {command.ExecutionId} is not in Running state (current: {execution.Status})");
        }

        var completedAt = DateTime.UtcNow;
        var executionTimeMs = (long)(completedAt - execution.StartedAt).TotalMilliseconds;

        execution.Output = command.Output;
        execution.Status = command.Status;
        execution.CompletedAt = completedAt;
        execution.ExecutionTimeMs = executionTimeMs;
        execution.InputTokens = command.InputTokens;
        execution.OutputTokens = command.OutputTokens;
        execution.TotalTokens = (command.InputTokens ?? 0) + (command.OutputTokens ?? 0);
        execution.EstimatedCost = command.EstimatedCost;
        execution.ToolCalls = command.ToolCalls;
        execution.ToolCallResults = command.ToolCallResults;
        execution.ErrorMessage = command.ErrorMessage;

        await _dbContext.SaveChangesAsync(cancellationToken);
    }
}

public class CompleteAgentExecutionCommandValidator : AbstractValidator<CompleteAgentExecutionCommand>
{
    public CompleteAgentExecutionCommandValidator()
    {
        RuleFor(x => x.ExecutionId)
            .NotEmpty().WithMessage("Execution ID is required");

        RuleFor(x => x.Status)
            .Must(s => s == ExecutionStatus.Completed || s == ExecutionStatus.Failed || s == ExecutionStatus.Cancelled)
            .WithMessage("Status must be Completed, Failed, or Cancelled");

        RuleFor(x => x.ErrorMessage)
            .NotEmpty()
            .When(x => x.Status == ExecutionStatus.Failed)
            .WithMessage("Error message is required when status is Failed");

        RuleFor(x => x.InputTokens)
            .GreaterThanOrEqualTo(0)
            .When(x => x.InputTokens.HasValue)
            .WithMessage("Input tokens must be >= 0");

        RuleFor(x => x.OutputTokens)
            .GreaterThanOrEqualTo(0)
            .When(x => x.OutputTokens.HasValue)
            .WithMessage("Output tokens must be >= 0");

        RuleFor(x => x.EstimatedCost)
            .GreaterThanOrEqualTo(0)
            .When(x => x.EstimatedCost.HasValue)
            .WithMessage("Estimated cost must be >= 0");
    }
}