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 NodeAdapterRegistry adapters = NodeAdapterRegistry.CreateDefault(
        new JobInboxStore(NodePaths.ForCurrentMachine().JobsDirectory));
    private readonly JobOutputUploader outputUploader = new(config);
    private readonly ConcurrentDictionary<Guid, Task> jobPipelines = new();
    private readonly ConcurrentDictionary<Guid, Task> exportPipelines = new();
    private readonly CancellationTokenSource forcedStop = new();
    private readonly WorkspaceStore workspaceStore = new(NodePaths.ForCurrentMachine());
    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)
    {
        try
        {
            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);
            }
        }
        finally
        {
            // A manual reconnect replaces this connection loop. Let any accepted job
            // finish its local pipeline first so the new loop cannot recover the same
            // worker concurrently and manufacture a conflicting terminal state.
            while (!jobPipelines.IsEmpty || !exportPipelines.IsEmpty)
            {
                await Task.WhenAll(jobPipelines.Values.Concat(exportPipelines.Values).ToArray());
            }
        }
    }

    internal void CancelActiveJobsForExit()
    {
        forcedStop.Cancel();
        foreach (var job in jobInbox.ReadRecoverableJobs().Where(item => item.Terminal is null))
        {
            adapters.Find(job.Capability)?.Cancel(job.JobId);
            if (!jobPipelines.ContainsKey(job.JobId))
            {
                WriteExitCancellation(job);
            }
        }
    }

    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
                    && !job.CloudTerminal)
                {
                    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)
                    {
                        if (!job.CloudTerminal) 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, adapters);
                    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))
                {
                    adapters.Find(cancelledJob.Capability)?.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);
                }
                else if (type.GetString() == "job_export"
                    && document.RootElement.TryGetProperty("payload", out var exportPayload))
                {
                    StartExport(exportPayload);
                }
            }
            message.SetLength(0);
        }
    }

    private void StartExport(JsonElement payload)
    {
        if (!payload.TryGetProperty("job_id", out var jobValue)
            || !Guid.TryParse(jobValue.GetString(), out var jobId)
            || !payload.TryGetProperty("lease_id", out var leaseValue)
            || !Guid.TryParse(leaseValue.GetString(), out var leaseId)
            || !payload.TryGetProperty("request_digest", out var digestValue)
            || !payload.TryGetProperty("output_ids", out var outputs)
            || outputs.ValueKind != JsonValueKind.Array)
        {
            return;
        }
        var job = jobInbox.ReadRecoverableJobs().SingleOrDefault(item =>
            item.JobId == jobId
            && item.LeaseId == leaseId
            && item.RequestDigest == digestValue.GetString()
            && item.Terminal is JsonElement terminal
            && terminal.GetProperty("status").GetString() == "succeeded");
        var outputIds = outputs.EnumerateArray()
            .Select(item => item.GetString())
            .Where(item => !string.IsNullOrWhiteSpace(item))
            .Cast<string>()
            .ToArray();
        if (job is null || outputIds.Length != outputs.GetArrayLength()) return;
        if (exportPipelines.TryAdd(jobId, Task.CompletedTask))
        {
            var running = RunExportAndReleaseAsync(job, outputIds);
            exportPipelines[jobId] = running;
        }
    }

    private async Task RunExportAndReleaseAsync(
        RecoverableJob job, IReadOnlyList<string> outputIds)
    {
        try
        {
            await outputUploader.ExportAsync(job, outputIds);
        }
        catch (Exception exception) when (
            exception is HttpRequestException
                or IOException
                or JsonException
                or UnauthorizedAccessException
                or InvalidDataException)
        {
            outputUploader.RecordDeferred(job, exception);
        }
        finally
        {
            exportPipelines.TryRemove(job.JobId, out _);
        }
    }

    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 adapter = adapters.Find(job.Capability)
            ?? throw new InvalidDataException($"No local adapter for {job.Capability}.");
        var recoveringOutputs = job.Terminal is JsonElement;
        try
        {
            if (job.Terminal is null)
            {
                jobInbox.WriteState(job, "downloading_inputs", 0, "正在下载并校验输入文件");
                await inputDownloader.DownloadAsync(job, forcedStop.Token);
                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, adapter.RunningDetail);
                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 adapter.RunAsync(job);
            }
            var refreshed = jobInbox.ReadRecoverableJobs()
                .Single(item => item.JobId == job.JobId);
            if (refreshed.Terminal is not JsonElement terminal)
            {
                throw new InvalidDataException("Software adapter did not create a terminal record.");
            }
            if (terminal.GetProperty("status").GetString() != "succeeded")
            {
                if (adapter.WorkspaceStateFilename is not null)
                {
                    workspaceStore.Restore(refreshed);
                }
                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;
            }
            if (adapter.WorkspaceStateFilename is not null)
            {
                workspaceStore.EnsureSucceeded(
                    refreshed,
                    adapter.Capability,
                    adapter.WorkspaceStateFilename);
            }
            jobInbox.WriteState(refreshed, "uploading_outputs", 90,
                adapter.WorkspaceStateFilename is null
                    ? "软件执行完成，正在上传结果"
                    : "软件执行完成，正在上传预览");
            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,
                adapter.WorkspaceStateFilename,
                adapter.PreviewOutputIds);
            if (jobInbox.ReadRecoverableJobs()
                .Single(item => item.JobId == job.JobId).CloudTerminal)
            {
                return;
            }
            jobInbox.WriteState(refreshed, "succeeded", 100,
                adapter.WorkspaceStateFilename is null
                    ? "结果已上传并由云端确认"
                    : "预览已上传，工程保存在本机工作区");
        }
        catch (OperationCanceledException) when (forcedStop.IsCancellationRequested)
        {
            var current = jobInbox.ReadRecoverableJobs()
                .Single(item => item.JobId == job.JobId);
            if (current.Terminal is null)
            {
                WriteExitCancellation(current);
            }
            if (jobInbox.ReadRecoverableJobs()
                .Single(item => item.JobId == job.JobId).Terminal is JsonElement terminal)
            {
                await TrySendAsync(socket, "job_terminal", terminal);
            }
        }
        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 void WriteExitCancellation(RecoverableJob job)
    {
        const string detail = "Cancelled by the local administrator while exiting the node.";
        jobInbox.WriteTerminal(job, "cancelled", "NODE_EXIT_CANCELLED", detail);
        jobInbox.WriteState(job, "cancelled", 0, "本机管理员退出节点并取消了任务");
    }

    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(NodePaths.ForCurrentMachine().RootDirectory)
            ?? throw new NodeConfigurationException("Data root has no drive.");
        var capabilityRuntime = adapters.All
            .ToDictionary(
                item => item.Capability,
                item =>
                {
                    var runtime = item.DetectRuntime();
                    var slots = runtime.Health == "ready"
                        && !jobInbox.HasPendingJobs
                        && !item.HasActiveJobs ? 1 : 0;
                    return new
                    {
                        software = runtime.Software,
                        software_version = runtime.SoftwareVersion,
                        adapter_version = item.AdapterVersion,
                        contract_sha256 = item.ContractSha256,
                        workspace_protocol_version = item.WorkspaceStateFilename is null ? 0 : 1,
                        features = item.Features,
                        available_slots = slots,
                        health = runtime.Health,
                        detail = runtime.Detail,
                    };
                });
        var installedCapabilities = adapters.All.Select(item => item.Capability).ToArray();
        // Keep the legacy top-level Origin payload during its deprecation window.
        var originAdapter = adapters.Find("origin.plot@v2");
        var origin = originAdapter?.DetectRuntime()
            ?? new AdapterRuntimeStatus(
                "OriginPro", null, "0.0.0", "unavailable", "Origin adapter is not installed.");
        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 = installedCapabilities,
            installed_capabilities = installedCapabilities,
            available_slots = origin.Health == "ready"
                && !jobInbox.HasPendingJobs
                && !(originAdapter?.HasActiveJobs ?? false) ? 1 : 0,
            capability_runtime = capabilityRuntime,
            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,
            },
        };
    }
}
