using System.ComponentModel;
using ClaudeDo.Data;
using Microsoft.EntityFrameworkCore;
using ModelContextProtocol.Server;
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;

namespace ClaudeDo.Worker.External;

public sealed record TaskStatusChangeDto(string TaskId, string Status);
public sealed record WaitForTaskChangeResult(IReadOnlyList<TaskStatusChangeDto> Changed, bool TimedOut);

[McpServerToolType]
public sealed class TaskWaitMcpTools
{
    // InteractiveLaunchSpecService sets MCP_TOOL_TIMEOUT=200000ms for the list handler
    // session; this cap leaves a ~30s margin so the tool itself reports TimedOut instead
    // of racing the client's own abort.
    internal const int MaxTimeoutSeconds = 170;
    private static readonly TimeSpan PollInterval = TimeSpan.FromMilliseconds(500);

    private readonly IDbContextFactory<ClaudeDoDbContext> _dbFactory;

    public TaskWaitMcpTools(IDbContextFactory<ClaudeDoDbContext> dbFactory)
    {
        _dbFactory = dbFactory;
    }

    [McpServerTool, Description(
        "Blocks until at least one of the given tasks leaves Queued/Running, or until timeoutSeconds elapses " +
        "(clamped server-side to 170s). Returns immediately if any task is already outside Queued/Running " +
        "when called (an unknown id is reported as status \"NotFound\" and counts as changed). Use this instead " +
        "of polling get_task in a loop. Result: { changed: [{ taskId, status }], timedOut }.")]
    public async Task<WaitForTaskChangeResult> WaitForTaskChange(
        string[] taskIds, int timeoutSeconds = 60, CancellationToken cancellationToken = default)
    {
        if (taskIds.Length == 0)
            throw new ArgumentException("taskIds must not be empty.", nameof(taskIds));

        var timeout = TimeSpan.FromSeconds(Math.Clamp(timeoutSeconds, 1, MaxTimeoutSeconds));
        using var timeoutCts = new CancellationTokenSource(timeout);
        using var linked = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCts.Token);

        try
        {
            while (true)
            {
                var changed = await CheckOnceAsync(taskIds, linked.Token);
                if (changed.Count > 0)
                    return new WaitForTaskChangeResult(changed, TimedOut: false);

                await Task.Delay(PollInterval, linked.Token);
            }
        }
        catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
        {
            return new WaitForTaskChangeResult(Array.Empty<TaskStatusChangeDto>(), TimedOut: true);
        }
    }

    private async Task<IReadOnlyList<TaskStatusChangeDto>> CheckOnceAsync(string[] taskIds, CancellationToken ct)
    {
        await using var ctx = await _dbFactory.CreateDbContextAsync(ct);
        var rows = await ctx.Tasks
            .AsNoTracking()
            .Where(t => taskIds.Contains(t.Id))
            .Select(t => new { t.Id, t.Status })
            .ToListAsync(ct);

        var byId = rows.ToDictionary(r => r.Id, r => r.Status);
        var result = new List<TaskStatusChangeDto>();
        foreach (var id in taskIds)
        {
            if (!byId.TryGetValue(id, out var status))
            {
                result.Add(new TaskStatusChangeDto(id, "NotFound"));
                continue;
            }
            if (status != TaskStatus.Queued && status != TaskStatus.Running)
                result.Add(new TaskStatusChangeDto(id, status.ToString()));
        }
        return result;
    }
}