using System.Collections.Concurrent;
using System.Threading.Channels;
using Agent.Core.Sessions;

namespace Agent.Core.Tasks;

public sealed class TaskRegistry
{
    private readonly ConcurrentDictionary<Guid, MutableTask> _tasks = new();
    private readonly SessionStore _sessions;
    private readonly Channel<Guid> _queue = Channel.CreateUnbounded<Guid>(new UnboundedChannelOptions { SingleReader = true, SingleWriter = false });

    public TaskRegistry(SessionStore? sessions = null) => _sessions = sessions ?? new SessionStore();

    public SessionStore Sessions => _sessions;

    public TaskSnapshot Enqueue(string request, Guid? requestedSessionId, double? temperature, IReadOnlyList<string>? images = null, Guid? requestedTaskId = null)
    {
        if (string.IsNullOrWhiteSpace(request)) throw new ArgumentException("Request cannot be empty.", nameof(request));
        if (temperature is < 0 or > 2) throw new ArgumentOutOfRangeException(nameof(temperature));
        Guid sessionId;
        if (requestedSessionId is { } existing)
        {
            if (!_sessions.TryGet(existing, out _)) throw new KeyNotFoundException($"Session '{existing}' was not found.");
            sessionId = existing;
        }
        else
        {
            sessionId = _sessions.Create().SessionId;
        }

        var task = new MutableTask(requestedTaskId ?? Guid.NewGuid(), sessionId, request, temperature, images ?? []);
        if (!_tasks.TryAdd(task.TaskId, task)) throw new InvalidOperationException("Could not register task.");
        if (!_queue.Writer.TryWrite(task.TaskId))
        {
            _tasks.TryRemove(task.TaskId, out _);
            throw new InvalidOperationException("Task queue is not accepting work.");
        }
        return task.Snapshot();
    }

    public bool TryGet(Guid taskId, Guid? sessionId, out TaskSnapshot? snapshot)
    {
        if (_tasks.TryGetValue(taskId, out MutableTask? task) && (sessionId is null || task.SessionId == sessionId))
        {
            snapshot = task.Snapshot();
            return true;
        }
        snapshot = null;
        return false;
    }

    public IAsyncEnumerable<Guid> ReadQueuedAsync(CancellationToken cancellationToken) => _queue.Reader.ReadAllAsync(cancellationToken);

    public bool Transition(Guid taskId, AgentTaskState state, string? currentTask = null, string? result = null)
    {
        return _tasks.TryGetValue(taskId, out MutableTask? task) && task.Transition(state, currentTask, result);
    }

    public bool AddUsage(Guid taskId, long inputTokens, long outputTokens) =>
        _tasks.TryGetValue(taskId, out MutableTask? task) && task.AddUsage(inputTokens, outputTokens);

    public CancellationToken GetCancellationToken(Guid taskId) =>
        _tasks.TryGetValue(taskId, out MutableTask? task) ? task.CancellationToken : CancellationToken.None;

    public bool Cancel(Guid taskId) =>
        _tasks.TryGetValue(taskId, out MutableTask? task) && task.Cancel();

    public async Task<TaskSnapshot> WaitForTerminalAsync(Guid taskId, CancellationToken cancellationToken)
    {
        if (!_tasks.TryGetValue(taskId, out MutableTask? task)) throw new KeyNotFoundException($"Task '{taskId}' was not found.");
        await task.Terminal.WaitAsync(cancellationToken).ConfigureAwait(false);
        return task.Snapshot();
    }

    private sealed class MutableTask(Guid taskId, Guid sessionId, string request, double? temperature, IReadOnlyList<string> images)
    {
        private readonly object _sync = new();
        private AgentTaskState _state = AgentTaskState.Queued;
        private string? _currentTask;
        private string? _result;
        private long _inputTokens;
        private long _outputTokens;
        private readonly CancellationTokenSource _cancellation = new();
        private readonly TaskCompletionSource _terminal = new(TaskCreationOptions.RunContinuationsAsynchronously);

        public Guid TaskId { get; } = taskId;
        public Guid SessionId { get; } = sessionId;
        public CancellationToken CancellationToken => _cancellation.Token;
        public Task Terminal => _terminal.Task;

        public TaskSnapshot Snapshot()
        {
            lock (_sync)
            {
                int completed = _state == AgentTaskState.Completed ? 1 : 0;
                return new(TaskId, SessionId, _state, request, temperature, _currentTask, _result, completed, 1, _inputTokens, _outputTokens, 0, images);
            }
        }

        public bool AddUsage(long inputTokens, long outputTokens)
        {
            if (inputTokens < 0 || outputTokens < 0) return false;
            lock (_sync)
            {
                _inputTokens += inputTokens;
                _outputTokens += outputTokens;
                return true;
            }
        }

        public bool Transition(AgentTaskState next, string? currentTask, string? result)
        {
            lock (_sync)
            {
                if (!Allowed(_state, next)) return false;
                _state = next;
                _currentTask = currentTask;
                _result = result;
                if (next is AgentTaskState.Completed or AgentTaskState.Failed or AgentTaskState.Cancelled) _terminal.TrySetResult();
                return true;
            }
        }

        public bool Cancel()
        {
            lock (_sync)
            {
                if (_state is AgentTaskState.Completed or AgentTaskState.Failed or AgentTaskState.Cancelled) return false;
                _cancellation.Cancel();
                _state = AgentTaskState.Cancelled;
                _currentTask = null;
                _result = "cancelled";
                _terminal.TrySetResult();
                return true;
            }
        }

        private static bool Allowed(AgentTaskState current, AgentTaskState next) => current switch
        {
            AgentTaskState.Queued => next is AgentTaskState.Planning or AgentTaskState.Failed or AgentTaskState.Cancelled,
            AgentTaskState.Planning => next is AgentTaskState.Running or AgentTaskState.Failed or AgentTaskState.Cancelled,
            AgentTaskState.Running => next is AgentTaskState.Planning or AgentTaskState.Completed or AgentTaskState.Failed or AgentTaskState.Cancelled,
            _ => false
        };
    }
}