using System.Collections.Concurrent;
using System.Net.WebSockets;
using System.Text;
using System.Text.Json;
using System.Text.Json.Nodes;
using backend.Realtime;
using backend.Tools;

namespace backend.DeviceHub;

public interface IDeviceClient
{
    Task SendEnvelopeAsync<T>(T envelope, CancellationToken ct) where T : HubEnvelope;
    Task SendBinaryAsync(ReadOnlyMemory<byte> bytes, CancellationToken ct);
}

public class ActiveDevice : IDeviceChannel, IDeviceClient, ISessionOutput, IAsyncDisposable
{
    public Guid DeviceId { get; }
    public WebSocket Ws => _ws;
    public RealtimeSession? CurrentSession { get; set; }
    public CancellationTokenSource? SessionCts { get; set; }

    private readonly WebSocket _ws;
    private readonly SemaphoreSlim _sendLock = new(1, 1);
    private readonly ConcurrentDictionary<string, TaskCompletionSource<JsonElement>> _pendingTools = new();

    public ActiveDevice(Guid deviceId, WebSocket ws)
    {
        DeviceId = deviceId;
        _ws = ws;
    }

    public async Task SendEnvelopeAsync<T>(T envelope, CancellationToken ct) where T : HubEnvelope
    {
        var json = JsonSerializer.Serialize(envelope, envelope.GetType());
        await _sendLock.WaitAsync(ct);
        try
        {
            await _ws.SendAsync(Encoding.UTF8.GetBytes(json),
                WebSocketMessageType.Text, endOfMessage: true, ct);
        }
        finally { _sendLock.Release(); }
    }

    public async Task SendBinaryAsync(ReadOnlyMemory<byte> bytes, CancellationToken ct)
    {
        await _sendLock.WaitAsync(ct);
        try
        {
            await _ws.SendAsync(bytes, WebSocketMessageType.Binary, endOfMessage: true, ct);
        }
        finally { _sendLock.Release(); }
    }

    public Task WriteEnvelopeAsync(JsonObject envelope, CancellationToken ct)
    {
        var json = envelope.ToJsonString();
        return WriteRawTextAsync(json, ct);
    }

    private async Task WriteRawTextAsync(string json, CancellationToken ct)
    {
        await _sendLock.WaitAsync(ct);
        try
        {
            await _ws.SendAsync(Encoding.UTF8.GetBytes(json),
                WebSocketMessageType.Text, endOfMessage: true, ct);
        }
        finally { _sendLock.Release(); }
    }

    public Task WriteBinaryAsync(ReadOnlyMemory<byte> bytes, CancellationToken ct)
        => SendBinaryAsync(bytes, ct);

    public Task<JsonElement> CallPiToolAsync(string name, JsonElement args, CancellationToken ct)
    {
        var callId = "srv_" + Guid.NewGuid().ToString("N")[..12];
        var tcs = new TaskCompletionSource<JsonElement>(TaskCreationOptions.RunContinuationsAsynchronously);
        _pendingTools[callId] = tcs;

        var envelope = new ToolCallEnvelope("tool_call", callId, name, args);
        _ = SendEnvelopeAsync(envelope, ct);

        ct.Register(() => tcs.TrySetCanceled());
        return tcs.Task;
    }

    public void CompleteToolResult(ToolResultEnvelope env)
    {
        if (!_pendingTools.TryRemove(env.CallId, out var tcs)) return;
        if (env.Ok && env.Result is JsonElement r) tcs.TrySetResult(r);
        else
        {
            var msg = (env.Error ?? "unknown").Replace("\\", "\\\\").Replace("\"", "\\\"");
            var doc = JsonDocument.Parse($$"""{"ok":false,"error":"{{msg}}"}""");
            tcs.TrySetResult(doc.RootElement.Clone());
        }
    }

    public async ValueTask DisposeAsync()
    {
        try
        {
            if (_ws.State == WebSocketState.Open)
                await _ws.CloseAsync(WebSocketCloseStatus.NormalClosure, "bye", CancellationToken.None);
        }
        catch { }
        _ws.Dispose();
        _sendLock.Dispose();
    }
}