using System.Collections.Concurrent;
using System.Text.Json;
using ModelContextProtocol.Extensions.Tasks;
using ModelContextProtocol.Protocol;

namespace Agent.Mcp;

public sealed record McpTaskStoreSettings(string Directory, TimeSpan TimeToLive, long PollIntervalMs);

public sealed class PersistentMcpTaskStore : IMcpTaskStore
{
    private readonly ConcurrentDictionary<CancellationToken, Guid> _taskIdsByCancellation = new();
    private readonly ConcurrentDictionary<string, StoredTask> _tasks = new(StringComparer.Ordinal);
    private readonly McpTaskStoreSettings _settings;
    private readonly JsonSerializerOptions _json = new(JsonSerializerDefaults.Web) { WriteIndented = true };

    public PersistentMcpTaskStore(McpTaskStoreSettings settings)
    {
        _settings = settings;
        System.IO.Directory.CreateDirectory(settings.Directory);
        Recover();
    }

    public event Action<InputResponseReceivedEventArgs>? InputResponseReceived;

    public Guid? ResolveTaskId(CancellationToken cancellationToken)
    {
        if (!_taskIdsByCancellation.TryRemove(cancellationToken, out Guid taskId)) return null;
        return taskId;
    }

    public async Task<McpTaskInfo> CreateTaskAsync(CancellationToken cancellationToken)
    {
        DateTimeOffset now = DateTimeOffset.UtcNow;
        var task = new StoredTask(Guid.NewGuid().ToString("D"), McpTaskStatus.Working, now, now,
            _settings.TimeToLive, _settings.PollIntervalMs, "working", null, null, null);
        if (!_tasks.TryAdd(task.TaskId, task)) throw new InvalidOperationException("Could not allocate an MCP task ID.");
        await SaveAsync(task, cancellationToken).ConfigureAwait(false);
        Guid createdId = Guid.Parse(task.TaskId);
        _taskIdsByCancellation[cancellationToken] = createdId;
        return ToInfo(task);
    }

    public Task<McpTaskInfo?> GetTaskAsync(string taskId, CancellationToken cancellationToken)
    {
        cancellationToken.ThrowIfCancellationRequested();
        if (!_tasks.TryGetValue(taskId, out StoredTask? task)) return Task.FromResult<McpTaskInfo?>(null);
        if (IsExpired(task))
        {
            _tasks.TryRemove(taskId, out _);
            Delete(taskId);
            return Task.FromResult<McpTaskInfo?>(null);
        }
        return Task.FromResult<McpTaskInfo?>(ToInfo(task));
    }

    public Task SetCompletedAsync(string taskId, JsonElement result, CancellationToken cancellationToken) =>
        UpdateAsync(taskId, McpTaskStatus.Completed, "completed", result, null, null, cancellationToken);

    public Task SetFailedAsync(string taskId, JsonElement error, CancellationToken cancellationToken) =>
        UpdateAsync(taskId, McpTaskStatus.Failed, "failed", null, error, null, cancellationToken);

    public async Task<bool> SetCancelledAsync(string taskId, CancellationToken cancellationToken)
    {
        if (!_tasks.ContainsKey(taskId)) return false;
        await UpdateAsync(taskId, McpTaskStatus.Cancelled, "cancelled", null, null, null, cancellationToken).ConfigureAwait(false);
        return true;
    }

    public async Task ResolveInputRequestsAsync(string taskId, IDictionary<string, InputResponse> responses, CancellationToken cancellationToken)
    {
        if (!_tasks.TryGetValue(taskId, out StoredTask? task)) return;
        var pending = task.InputRequests is null
            ? new Dictionary<string, InputRequest>(StringComparer.Ordinal)
            : new Dictionary<string, InputRequest>(task.InputRequests, StringComparer.Ordinal);
        var accepted = new List<KeyValuePair<string, InputResponse>>();
        foreach ((string requestId, InputResponse response) in responses)
        {
            if (!pending.Remove(requestId)) continue;
            accepted.Add(new KeyValuePair<string, InputResponse>(requestId, response));
        }
        await UpdateAsync(taskId, pending.Count == 0 ? McpTaskStatus.Working : McpTaskStatus.InputRequired,
            pending.Count == 0 ? "working" : "input required", task.Result, task.Error,
            pending.Count == 0 ? null : pending, cancellationToken).ConfigureAwait(false);
        foreach ((string requestId, InputResponse response) in accepted)
            InputResponseReceived?.Invoke(new InputResponseReceivedEventArgs { TaskId = taskId, RequestId = requestId, Response = response });
    }

    public Task SetInputRequestsAsync(string taskId, IDictionary<string, InputRequest> requests, CancellationToken cancellationToken) =>
        UpdateAsync(taskId, McpTaskStatus.InputRequired, "input required", null, null,
            new Dictionary<string, InputRequest>(requests, StringComparer.Ordinal), cancellationToken);

    private async Task UpdateAsync(string taskId, McpTaskStatus status, string message, JsonElement? result,
        JsonElement? error, IReadOnlyDictionary<string, InputRequest>? inputRequests, CancellationToken cancellationToken)
    {
        if (!_tasks.TryGetValue(taskId, out StoredTask? current)) throw new KeyNotFoundException($"MCP task '{taskId}' was not found.");
        var updated = current with
        {
            Status = status,
            LastUpdatedAt = DateTimeOffset.UtcNow,
            StatusMessage = message,
            Result = result,
            Error = error,
            InputRequests = inputRequests
        };
        _tasks[taskId] = updated;
        await SaveAsync(updated, cancellationToken).ConfigureAwait(false);
    }

    private void Recover()
    {
        foreach (string path in System.IO.Directory.EnumerateFiles(_settings.Directory, "*.json", SearchOption.TopDirectoryOnly))
        {
            try
            {
                StoredTask? task = JsonSerializer.Deserialize<StoredTask>(File.ReadAllText(path), _json);
                if (task is null || IsExpired(task)) { File.Delete(path); continue; }
                if (task.Status is McpTaskStatus.Working or McpTaskStatus.InputRequired)
                {
                    JsonElement error = JsonSerializer.SerializeToElement(new { code = "server_restarted", message = "The server restarted before the task completed." });
                    task = task with { Status = McpTaskStatus.Failed, StatusMessage = "server_restarted", Error = error, LastUpdatedAt = DateTimeOffset.UtcNow, InputRequests = null };
                    SaveSynchronously(task);
                }
                _tasks[task.TaskId] = task;
            }
            catch (JsonException)
            {
                File.Move(path, path + ".invalid", true);
            }
        }
    }

    private bool IsExpired(StoredTask task) => task.TimeToLive is { } ttl && DateTimeOffset.UtcNow - task.CreatedAt >= ttl;

    private async ValueTask SaveAsync(StoredTask task, CancellationToken cancellationToken)
    {
        string path = PathFor(task.TaskId);
        string temp = path + ".tmp-" + Guid.NewGuid().ToString("N");
        try
        {
            await using (var stream = new FileStream(temp, FileMode.CreateNew, FileAccess.Write, FileShare.None, 4096,
                FileOptions.Asynchronous | FileOptions.WriteThrough))
            {
                await JsonSerializer.SerializeAsync(stream, task, _json, cancellationToken).ConfigureAwait(false);
                await stream.FlushAsync(cancellationToken).ConfigureAwait(false);
            }
            File.Move(temp, path, true);
        }
        finally { if (File.Exists(temp)) File.Delete(temp); }
    }

    private void SaveSynchronously(StoredTask task)
    {
        string path = PathFor(task.TaskId);
        string temp = path + ".tmp-" + Guid.NewGuid().ToString("N");
        try
        {
            using (var stream = new FileStream(temp, FileMode.CreateNew, FileAccess.Write, FileShare.None, 4096, FileOptions.WriteThrough))
            {
                JsonSerializer.Serialize(stream, task, _json);
                stream.Flush(true);
            }
            File.Move(temp, path, true);
        }
        finally { if (File.Exists(temp)) File.Delete(temp); }
    }

    private void Delete(string taskId)
    {
        string path = PathFor(taskId);
        if (File.Exists(path)) File.Delete(path);
    }

    private string PathFor(string taskId) => Path.Combine(_settings.Directory, taskId + ".json");

    private static McpTaskInfo ToInfo(StoredTask task) => new(task.TaskId, task.Status, task.CreatedAt, task.LastUpdatedAt,
        task.TimeToLive, task.PollIntervalMs, task.StatusMessage, task.Result, task.Error, task.InputRequests);

    private sealed record StoredTask(
        string TaskId,
        McpTaskStatus Status,
        DateTimeOffset CreatedAt,
        DateTimeOffset LastUpdatedAt,
        TimeSpan? TimeToLive,
        long? PollIntervalMs,
        string StatusMessage,
        JsonElement? Result,
        JsonElement? Error,
        IReadOnlyDictionary<string, InputRequest>? InputRequests);
}