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

namespace Zcbot.WindowsNode;

internal sealed class JobCoordinator : IDisposable
{
    private readonly NodeConfig config;
    private readonly NodePaths paths;
    private readonly JobInboxStore offerInbox;
    private readonly JobRepository jobInbox;
    private readonly JobRecoveryService jobRecovery;
    private readonly NodeHttpClient httpClient;
    private readonly JobInputDownloader inputDownloader;
    private readonly AdapterCatalog adapters;
    private readonly JobOutputUploader outputUploader;
    private readonly ConcurrentDictionary<Guid, Task> jobPipelines = new();
    private readonly ConcurrentDictionary<Guid, Task> exportPipelines = new();
    private readonly CancellationTokenSource forcedStop = new();
    private readonly WorkspaceStore workspaceStore;

    internal JobCoordinator(NodeConfig config, NodePaths paths)
    {
        this.config = config;
        this.paths = paths;
        offerInbox = new JobInboxStore(paths.JobsDirectory);
        jobInbox = new JobRepository(paths.JobsDirectory);
        jobRecovery = new JobRecoveryService(jobInbox);
        httpClient = new NodeHttpClient(config);
        inputDownloader = new JobInputDownloader(httpClient, jobInbox);
        adapters = AdapterCatalog.CreateDefault(jobInbox, paths);
        outputUploader = new JobOutputUploader(httpClient, paths);
        workspaceStore = new WorkspaceStore(paths);
    }

    internal async Task WaitForIdleAsync()
    {
        // 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);
            }
        }
    }

    internal Task ResumeDeferredUploadsAsync(
        NodeProtocolClient client, CancellationToken cancellationToken)
    {
        foreach (var decision in jobRecovery.Scan())
        {
            if (decision.Action == JobRecoveryAction.ResumeOutputUpload)
            {
                StartJobPipeline(client, decision.Job, recovery: true);
            }
        }
        return Task.CompletedTask;
    }

    internal async Task ReportRecoverableJobsAsync(
        NodeProtocolClient client, CancellationToken cancellationToken)
    {
        foreach (var decision in jobRecovery.Scan())
        {
            var job = decision.Job;
            if (decision.Action == JobRecoveryAction.ReplayTerminal)
            {
                await client.SendAsync("job_terminal", job.Terminal!.Value, cancellationToken);
                continue;
            }
            if (decision.Action == JobRecoveryAction.ResumeOutputUpload)
            {
                StartJobPipeline(client, job, recovery: true);
                continue;
            }
            if (decision.Action == JobRecoveryAction.None)
            {
                continue;
            }
            await client.SendAsync("job_state", new
            {
                job_id = job.JobId,
                lease_id = job.LeaseId,
                request_digest = job.RequestDigest,
                stage = "downloading_inputs",
                progress = 0,
                metrics = new { },
            }, cancellationToken);
            StartJobPipeline(client, job, recovery: true);
        }
    }

    internal async Task HandleMessageAsync(
        NodeProtocolClient client,
        NodeProtocolMessage message,
        CancellationToken cancellationToken)
    {
        var payload = message.Payload;
        if (message.Type == "job_offer" && payload.ValueKind != JsonValueKind.Undefined)
        {
            var offerResult = offerInbox.Accept(payload, adapters);
            await client.SendAsync(
                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 client.SendAsync("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(client, acceptedJob);
            }
        }
        else if (message.Type == "job_cancel"
            && payload.ValueKind != JsonValueKind.Undefined
            && TryCancelJob(payload, out var cancelledJob))
        {
            adapters.Find(cancelledJob.Capability)?.Cancel(cancelledJob.JobId);
            jobInbox.WriteTerminal(
                cancelledJob, "cancelled", "USER_CANCELLED", "Cancelled by user.");
            await client.SendAsync("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 (message.Type == "job_export"
            && payload.ValueKind != JsonValueKind.Undefined)
        {
            StartExport(payload);
        }
    }

    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(
        NodeProtocolClient socket,
        RecoverableJob job,
        bool recovery = false)
    {
        var completion = new TaskCompletionSource(
            TaskCreationOptions.RunContinuationsAsynchronously);
        if (jobPipelines.TryAdd(job.JobId, completion.Task))
        {
            _ = RunJobPipelineAndReleaseAsync(socket, job, recovery, completion);
        }
    }

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

    private async Task RunJobPipelineAsync(
        NodeProtocolClient socket,
        RecoverableJob job,
        bool recovery)
    {
        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,
                    "正在下载并校验输入文件",
                    recovery ? JobTransitionMode.Recovery : JobTransitionMode.Normal);
                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
                    ? "软件执行完成，正在上传结果"
                    : "软件执行完成，正在上传预览",
                recovery ? JobTransitionMode.Recovery : JobTransitionMode.Normal);
            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(NodeProtocolClient 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(NodeProtocolClient socket, string type, object payload)
    {
        if (!socket.IsOpen)
        {
            return;
        }
        try
        {
            await socket.SendAsync(type, payload, CancellationToken.None);
        }
        catch (Exception exception) when (
            exception is WebSocketException or IOException or ObjectDisposedException)
        {
            Console.Error.WriteLine($"[WARN] Job report deferred: {exception.Message}");
        }
    }

    internal object RuntimePayload()
    {
        var root = Path.GetPathRoot(paths.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,
            },
        };
    }

    public void Dispose()
    {
        forcedStop.Dispose();
        httpClient.Dispose();
    }
}