using System.Net;
using System.Net.WebSockets;
using System.Reflection;
using System.Runtime.InteropServices;
using System.Text;
using System.Text.Json;
using System.Collections.Concurrent;

namespace Zcbot.WindowsNode;

internal sealed class NodeConnectionLoop(NodeConfig config, Action<NodeStatus>? statusChanged = null)
{
    private readonly SemaphoreSlim sendLock = new(1, 1);
    private readonly JobInboxStore jobInbox = new(NodePaths.ForCurrentMachine().JobsDirectory);
    private readonly JobInputDownloader inputDownloader = new(
        config, new JobInboxStore(NodePaths.ForCurrentMachine().JobsDirectory));
    private readonly OriginWorkerRunner workerRunner = new(
        new JobInboxStore(NodePaths.ForCurrentMachine().JobsDirectory));
    private readonly JobOutputUploader outputUploader = new(config);
    private readonly ConcurrentDictionary<Guid, Task> jobPipelines = new();
    private static readonly TimeSpan[] Backoff =
    [
        TimeSpan.FromSeconds(1),
        TimeSpan.FromSeconds(2),
        TimeSpan.FromSeconds(5),
        TimeSpan.FromSeconds(10),
        TimeSpan.FromSeconds(30),
        TimeSpan.FromSeconds(60),
    ];

    internal async Task RunAsync(CancellationToken cancellationToken)
    {
        var attempt = 0;
        while (!cancellationToken.IsCancellationRequested)
        {
            try
            {
                Report(NodeState.Connecting, "正在连接 zcbot…");
                await ConnectOnceAsync(cancellationToken);
                attempt = 0;
            }
            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
            {
                throw;
            }
            catch (WebSocketException exception)
            {
                Report(NodeState.Offline, "连接中断，等待重连");
                Console.Error.WriteLine($"[WARN] WebSocket disconnected: {exception.Message}");
            }
            catch (IOException exception)
            {
                Report(NodeState.Offline, "网络不可用，等待重连");
                Console.Error.WriteLine($"[WARN] Connection I/O failed: {exception.Message}");
            }
            catch (JsonException exception)
            {
                Report(NodeState.Offline, "服务端消息无效，等待重连");
                Console.Error.WriteLine($"[WARN] Invalid server message: {exception.Message}");
            }
            catch (NodeEndpointException exception)
            {
                Report(NodeState.Offline, "WebSocket 握手被拒绝，请检查服务端或反向代理");
                Console.Error.WriteLine($"[WARN] WebSocket handshake rejected: {exception.Message}");
            }

            var baseDelay = Backoff[Math.Min(attempt, Backoff.Length - 1)];
            attempt++;
            var jitter = TimeSpan.FromMilliseconds(Random.Shared.Next(0, 750));
            var delay = baseDelay + jitter;
            Console.WriteLine($"[INFO] Reconnecting in {delay.TotalSeconds:F1}s.");
            await Task.Delay(delay, cancellationToken);
        }
    }

    private async Task ConnectOnceAsync(CancellationToken cancellationToken)
    {
        using var socket = new ClientWebSocket();
        socket.Options.SetRequestHeader("Authorization", $"Bearer {config.NodeToken}");
        socket.Options.SetRequestHeader("X-Node-Id", config.NodeId.ToString());
        socket.Options.KeepAliveInterval = TimeSpan.FromSeconds(config.HeartbeatSeconds);
        socket.Options.CollectHttpResponseDetails = true;

        var endpoint = NodeUri.WebSocketEndpoint(config.ServerUrl);
        Console.WriteLine($"[INFO] Connecting to {endpoint.GetLeftPart(UriPartial.Path)}.");
        try
        {
            await socket.ConnectAsync(endpoint, cancellationToken);
        }
        catch (WebSocketException exception) when (
            socket.HttpStatusCode is HttpStatusCode.Unauthorized or HttpStatusCode.Forbidden)
        {
            throw new NodeEndpointException(
                $"WebSocket handshake returned HTTP {(int?)socket.HttpStatusCode}. "
                + "Verify that the server includes the current Windows Node routes and that "
                + $"the reverse proxy forwards WebSocket Upgrade for {endpoint.AbsolutePath}. "
                + exception.Message);
        }
        Console.WriteLine("[OK] Node connected.");
        Report(NodeState.Online, "已连接");

        await SendAsync(socket, "hello", RuntimePayload(), cancellationToken);
        await ReportRecoverableJobsAsync(socket, cancellationToken);
        using var heartbeatStop = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        var heartbeat = HeartbeatLoopAsync(socket, heartbeatStop.Token);
        try
        {
            await ReceiveLoopAsync(socket, cancellationToken);
        }
        finally
        {
            heartbeatStop.Cancel();
            try
            {
                await heartbeat;
            }
            catch (OperationCanceledException) when (heartbeatStop.IsCancellationRequested)
            {
            }
        }
    }

    private async Task HeartbeatLoopAsync(ClientWebSocket socket, CancellationToken cancellationToken)
    {
        using var timer = new PeriodicTimer(TimeSpan.FromSeconds(config.HeartbeatSeconds));
        while (await timer.WaitForNextTickAsync(cancellationToken))
        {
            await SendAsync(socket, "heartbeat", RuntimePayload(), cancellationToken);
            foreach (var job in jobInbox.ReadRecoverableJobs())
            {
                if (job.Terminal is JsonElement terminal
                    && terminal.GetProperty("status").GetString() == "succeeded"
                    && !job.UploadComplete)
                {
                    StartJobPipeline(socket, job);
                }
            }
        }
    }

    private async Task ReportRecoverableJobsAsync(
        ClientWebSocket socket, CancellationToken cancellationToken)
    {
        foreach (var job in jobInbox.ReadRecoverableJobs())
        {
            if (job.Terminal is JsonElement terminal)
            {
                if (terminal.GetProperty("status").GetString() == "succeeded")
                {
                    if (!job.UploadComplete)
                    {
                        StartJobPipeline(socket, job);
                    }
                }
                else
                {
                    await SendAsync(socket, "job_terminal", terminal, cancellationToken);
                }
                continue;
            }
            await SendAsync(socket, "job_state", new
            {
                job_id = job.JobId,
                lease_id = job.LeaseId,
                request_digest = job.RequestDigest,
                stage = "downloading_inputs",
                progress = 0,
                metrics = new { },
            }, cancellationToken);
            StartJobPipeline(socket, job);
        }
    }

    private async Task ReceiveLoopAsync(
        ClientWebSocket socket, CancellationToken cancellationToken)
    {
        var buffer = new byte[16 * 1024];
        using var message = new MemoryStream();
        while (socket.State == WebSocketState.Open)
        {
            var result = await socket.ReceiveAsync(buffer, cancellationToken);
            if (result.MessageType == WebSocketMessageType.Close)
            {
                if ((int?)result.CloseStatus == 4003)
                {
                    throw new NodeConfigurationException(
                        "节点身份已被服务端拒绝。请在管理后台确认该 Node ID 未被禁用或删除；"
                        + "若记录不存在或身份已撤销，请清除本机身份并使用新注册码重新注册。"
                        + $" 服务端信息：{result.CloseStatusDescription}");
                }
                if (result.CloseStatus == WebSocketCloseStatus.PolicyViolation)
                {
                    throw new NodeConfigurationException(
                        $"节点上报被服务端拒绝：{result.CloseStatusDescription}");
                }
                return;
            }
            if (result.MessageType != WebSocketMessageType.Text)
            {
                throw new JsonException("Only text WebSocket messages are supported.");
            }
            message.Write(buffer, 0, result.Count);
            if (!result.EndOfMessage)
            {
                if (message.Length > 1024 * 1024)
                {
                    throw new JsonException("Server message exceeded 1 MiB.");
                }
                continue;
            }
            using var document = JsonDocument.Parse(message.ToArray());
            if (document.RootElement.TryGetProperty("type", out var type))
            {
                Console.WriteLine($"[INFO] Server message: {type.GetString()}.");
                if (type.GetString() == "job_offer"
                    && document.RootElement.TryGetProperty("payload", out var payload))
                {
                    var offerResult = jobInbox.Accept(payload);
                    await SendAsync(
                        socket,
                        offerResult.Accepted ? "job_accept" : "job_reject",
                        offerResult.Accepted
                            ? new
                            {
                                job_id = offerResult.JobId,
                                lease_id = offerResult.LeaseId,
                                request_digest = offerResult.RequestDigest,
                            }
                            : new
                            {
                                job_id = payload.TryGetProperty("job_id", out var jobId)
                                    ? jobId.GetString() : "",
                                lease_id = payload.TryGetProperty("lease_id", out var leaseId)
                                    ? leaseId.GetString() : "",
                                reason = offerResult.Reason,
                            },
                        cancellationToken);
                    if (offerResult.Accepted)
                    {
                        await SendAsync(socket, "job_state", new
                        {
                            job_id = offerResult.JobId,
                            lease_id = offerResult.LeaseId,
                            request_digest = offerResult.RequestDigest,
                            stage = "downloading_inputs",
                            progress = 0,
                            metrics = new { },
                        }, cancellationToken);
                        var acceptedJob = jobInbox.ReadRecoverableJobs()
                            .Single(item => item.JobId == offerResult.JobId);
                        StartJobPipeline(socket, acceptedJob);
                    }
                }
                else if (type.GetString() == "job_cancel"
                    && document.RootElement.TryGetProperty("payload", out var cancelPayload)
                    && TryCancelJob(cancelPayload, out var cancelledJob))
                {
                    workerRunner.Cancel(cancelledJob.JobId);
                    jobInbox.WriteTerminal(
                        cancelledJob, "cancelled", "USER_CANCELLED", "Cancelled by user.");
                    await SendAsync(socket, "job_terminal", new
                    {
                        job_id = cancelledJob.JobId,
                        lease_id = cancelledJob.LeaseId,
                        request_digest = cancelledJob.RequestDigest,
                        status = "cancelled",
                        error = new { code = "USER_CANCELLED", detail = "Cancelled by user." },
                        artifact_manifest = Array.Empty<object>(),
                    }, cancellationToken);
                }
            }
            message.SetLength(0);
        }
    }

    private bool TryCancelJob(JsonElement payload, out RecoverableJob job)
    {
        job = null!;
        if (!payload.TryGetProperty("job_id", out var jobIdValue)
            || !Guid.TryParse(jobIdValue.GetString(), out var jobId)
            || !payload.TryGetProperty("lease_id", out var leaseIdValue)
            || !Guid.TryParse(leaseIdValue.GetString(), out var leaseId)
            || !payload.TryGetProperty("request_digest", out var digestValue))
        {
            return false;
        }
        job = jobInbox.ReadRecoverableJobs().SingleOrDefault(item =>
            item.JobId == jobId
            && item.LeaseId == leaseId
            && item.RequestDigest == digestValue.GetString()
            && item.Terminal is null)!;
        return job is not null;
    }

    private void StartJobPipeline(ClientWebSocket socket, RecoverableJob job)
    {
        var completion = new TaskCompletionSource(
            TaskCreationOptions.RunContinuationsAsynchronously);
        if (jobPipelines.TryAdd(job.JobId, completion.Task))
        {
            _ = RunJobPipelineAndReleaseAsync(socket, job, completion);
        }
    }

    private async Task RunJobPipelineAndReleaseAsync(
        ClientWebSocket socket, RecoverableJob job, TaskCompletionSource completion)
    {
        try
        {
            await RunJobPipelineAsync(socket, job);
        }
        finally
        {
            completion.TrySetResult();
            jobPipelines.TryRemove(job.JobId, out _);
        }
    }

    private async Task RunJobPipelineAsync(ClientWebSocket socket, RecoverableJob job)
    {
        var recoveringOutputs = job.Terminal is JsonElement;
        try
        {
            if (job.Terminal is null)
            {
                jobInbox.WriteState(job, "downloading_inputs", 0, "正在下载并校验输入文件");
                await inputDownloader.DownloadAsync(job, CancellationToken.None);
                var afterDownload = jobInbox.ReadRecoverableJobs()
                    .Single(item => item.JobId == job.JobId);
                if (afterDownload.Terminal is JsonElement cancelledTerminal)
                {
                    await TrySendAsync(socket, "job_terminal", cancelledTerminal);
                    return;
                }
                jobInbox.WriteState(job, "ready_to_run", 5, "输入文件已就绪，准备启动软件");
                await TrySendAsync(socket, "job_state", new
                {
                    job_id = job.JobId,
                    lease_id = job.LeaseId,
                    request_digest = job.RequestDigest,
                    stage = "ready_to_run",
                    progress = 5,
                    metrics = new
                    {
                        input_count = jobInbox.InputTransfers(job).Count,
                        input_bytes = jobInbox.InputTransfers(job)
                            .Sum(item => item.GetProperty("size_bytes").GetInt64()),
                    },
                });
                jobInbox.WriteState(job, "software_running", 10, "Origin 正在生成图形");
                await TrySendAsync(socket, "job_state", new
                {
                    job_id = job.JobId,
                    lease_id = job.LeaseId,
                    request_digest = job.RequestDigest,
                    stage = "software_running",
                    progress = 10,
                    metrics = new { },
                });
                await workerRunner.RunAsync(job);
            }
            var refreshed = jobInbox.ReadRecoverableJobs()
                .Single(item => item.JobId == job.JobId);
            if (refreshed.Terminal is not JsonElement terminal)
            {
                throw new InvalidDataException("Origin worker did not create a terminal record.");
            }
            if (terminal.GetProperty("status").GetString() != "succeeded")
            {
                var terminalStatus = terminal.GetProperty("status").GetString() ?? "failed";
                var terminalDetail = terminal.TryGetProperty("error", out var error)
                    && error.TryGetProperty("detail", out var errorDetail)
                    ? errorDetail.GetString() ?? "任务执行失败"
                    : "任务执行失败";
                jobInbox.WriteState(
                    refreshed, terminalStatus, terminalStatus == "cancelled" ? 0 : 10, terminalDetail);
                await TrySendAsync(socket, "job_terminal", terminal);
                return;
            }
            jobInbox.WriteState(refreshed, "uploading_outputs", 90, "软件执行完成，正在上传结果");
            await TrySendAsync(socket, "job_state", new
            {
                job_id = job.JobId,
                lease_id = job.LeaseId,
                request_digest = job.RequestDigest,
                stage = "uploading_outputs",
                progress = 90,
                metrics = new { },
            });
            await outputUploader.UploadAsync(refreshed, recoveringOutputs);
            jobInbox.WriteState(refreshed, "succeeded", 100, "结果已上传并由云端确认");
        }
        catch (Exception exception) when (
            exception is HttpRequestException
                or IOException
                or JsonException
                or UnauthorizedAccessException
                or InvalidDataException)
        {
            var current = jobInbox.ReadRecoverableJobs()
                .Single(item => item.JobId == job.JobId);
            if (current.Terminal is JsonElement terminal
                && terminal.GetProperty("status").GetString() == "succeeded")
            {
                jobInbox.WriteState(
                    current,
                    "uploading_outputs",
                    90,
                    $"结果上传暂缓：{exception.Message}");
                outputUploader.RecordDeferred(current, exception);
                return;
            }
            jobInbox.WriteTerminal(
                job,
                "failed",
                "INPUT_DOWNLOAD_FAILED",
                exception.Message[..Math.Min(exception.Message.Length, 500)]);
            jobInbox.WriteState(job, "failed", 0, exception.Message);
            var failedTerminal = jobInbox.ReadRecoverableJobs()
                .Single(item => item.JobId == job.JobId).Terminal;
            if (failedTerminal is JsonElement payload)
            {
                await TrySendAsync(socket, "job_terminal", payload);
            }
        }
    }

    private async Task TrySendTerminalAsync(ClientWebSocket socket, RecoverableJob job)
    {
        var terminal = jobInbox.ReadRecoverableJobs()
            .Single(item => item.JobId == job.JobId).Terminal;
        if (terminal is JsonElement terminalPayload)
        {
            await TrySendAsync(socket, "job_terminal", terminalPayload);
        }
    }

    private async Task TrySendAsync(ClientWebSocket socket, string type, object payload)
    {
        if (socket.State != WebSocketState.Open)
        {
            return;
        }
        try
        {
            await SendAsync(socket, type, payload, CancellationToken.None);
        }
        catch (Exception exception) when (
            exception is WebSocketException or IOException or ObjectDisposedException)
        {
            Console.Error.WriteLine($"[WARN] Job report deferred: {exception.Message}");
        }
    }

    private void Report(NodeState state, string message) =>
        statusChanged?.Invoke(NodeStatus.Create(state, message));

    private async Task SendAsync(
        ClientWebSocket socket, string type, object payload, CancellationToken cancellationToken)
    {
        var envelope = JsonSerializer.SerializeToUtf8Bytes(new
        {
            protocol_version = 1,
            message_id = Guid.NewGuid(),
            type,
            sent_at = DateTimeOffset.UtcNow,
            payload,
        });
        await sendLock.WaitAsync(cancellationToken);
        try
        {
            await socket.SendAsync(
                envelope, WebSocketMessageType.Text, endOfMessage: true, cancellationToken);
        }
        finally
        {
            sendLock.Release();
        }
    }

    private object RuntimePayload()
    {
        var root = Path.GetPathRoot(Environment.SystemDirectory) ?? "C:\\";
        var origin = OriginRuntimeProbe.Detect();
        return new
        {
            install_id = config.InstallId,
            node_name = config.NodeName,
            node_version = Assembly.GetExecutingAssembly().GetName().Version?.ToString(3) ?? "0.1.0",
            os_version = RuntimeInformation.OSDescription,
            capabilities = config.Capabilities,
            available_slots = origin.Health == "ready"
                && !jobInbox.HasPendingOriginJobs
                && !workerRunner.HasActiveJobs ? 1 : 0,
            disk_free_bytes = new DriveInfo(root).AvailableFreeSpace,
            desktop_session = Environment.UserInteractive,
            origin = new
            {
                software = origin.Software,
                software_version = origin.SoftwareVersion,
                adapter_version = origin.AdapterVersion,
                health = origin.Health,
                detail = origin.Detail,
            },
        };
    }
}