using System.Text;
using System.Text.Json;

namespace Zcbot.WindowsNode;

internal sealed class JobInboxStore(string jobsDirectory)
{
    private static readonly JsonSerializerOptions JsonOptions = new() { WriteIndented = true };
    private static readonly HashSet<string> PlotTypes =
    ["line", "scatter", "line_scatter"];
    private static readonly HashSet<string> OutputFormats = ["opju", "png", "svg", "pdf"];

    internal bool HasPendingJobs => Directory.Exists(jobsDirectory)
        && ReadRecoverableJobs().Any(item =>
            item.Terminal is null
            || item.Terminal.Value.GetProperty("status").GetString() == "succeeded"
                && !item.UploadComplete);

    internal IReadOnlyList<RecoverableJob> ReadRecoverableJobs()
    {
        if (!Directory.Exists(jobsDirectory))
        {
            return [];
        }
        var jobs = new List<RecoverableJob>();
        foreach (var requestPath in Directory.EnumerateFiles(
            jobsDirectory, "request.json", SearchOption.AllDirectories))
        {
            try
            {
                using var request = JsonDocument.Parse(File.ReadAllBytes(requestPath));
                var root = request.RootElement;
                if (!TryReadGuid(root, "job_id", out var jobId)
                    || !TryReadGuid(root, "lease_id", out var leaseId)
                    || !root.TryGetProperty("request_digest", out var digestValue)
                    || digestValue.GetString() is not { Length: 64 } requestDigest)
                {
                    continue;
                }
                var jobDirectory = Directory.GetParent(Directory.GetParent(requestPath)!.FullName)!.FullName;
                jobs.Add(new RecoverableJob(
                    jobId,
                    leaseId,
                    requestDigest,
                    root.TryGetProperty("input_transfer", out var transfer)
                        ? transfer.Clone() : null,
                    ReadTerminal(Path.Combine(jobDirectory, "terminal.json")),
                    File.Exists(Path.Combine(jobDirectory, "upload-complete.json"))));
            }
            catch (Exception exception) when (
                exception is JsonException or IOException or UnauthorizedAccessException)
            {
            }
        }
        return jobs;
    }

    internal JobOfferResult Accept(JsonElement payload)
    {
        if (!TryReadGuid(payload, "job_id", out var jobId)
            || !TryReadGuid(payload, "lease_id", out var leaseId)
            || !payload.TryGetProperty("request_digest", out var digestValue)
            || digestValue.ValueKind != JsonValueKind.String
            || digestValue.GetString() is not { Length: 64 } requestDigest
            || !payload.TryGetProperty("capability", out var capabilityValue)
            || capabilityValue.GetString() != "origin.plot@v1"
            || !payload.TryGetProperty("request", out var request)
            || request.ValueKind != JsonValueKind.Object
            || !payload.TryGetProperty("input_transfer", out var inputTransfer)
            || !IsValidInputTransfer(inputTransfer))
        {
            return JobOfferResult.Reject("invalid_offer");
        }
        if (!IsValidRequest(request))
        {
            return JobOfferResult.Reject("unsupported_request");
        }

        var directory = Path.Combine(jobsDirectory, jobId.ToString("D"));
        var requestDirectory = Path.Combine(directory, "request");
        var requestPath = Path.Combine(requestDirectory, "request.json");
        Directory.CreateDirectory(requestDirectory);
        if (File.Exists(requestPath))
        {
            try
            {
                using var existing = JsonDocument.Parse(File.ReadAllBytes(requestPath));
                var root = existing.RootElement;
                var sameDigest = root.TryGetProperty("request_digest", out var existingDigest)
                    && existingDigest.GetString() == requestDigest;
                if (!sameDigest)
                {
                    return JobOfferResult.Reject("job_digest_conflict");
                }
                if (!TryReadGuid(root, "lease_id", out var existingLease)
                    || existingLease != leaseId)
                {
                    var updated = JsonSerializer.SerializeToUtf8Bytes(new
                    {
                        job_id = jobId,
                        lease_id = leaseId,
                        request_digest = requestDigest,
                        capability = "origin.plot@v1",
                        accepted_at = DateTimeOffset.UtcNow,
                        request,
                        input_transfer = inputTransfer,
                    }, JsonOptions);
                    AtomicWrite(requestPath, updated, overwrite: true);
                }
                return JobOfferResult.Accept(jobId, leaseId, requestDigest);
            }
            catch (JsonException)
            {
                return JobOfferResult.Reject("local_job_record_invalid");
            }
        }

        var record = JsonSerializer.SerializeToUtf8Bytes(new
        {
            job_id = jobId,
            lease_id = leaseId,
            request_digest = requestDigest,
            capability = "origin.plot@v1",
            accepted_at = DateTimeOffset.UtcNow,
            request,
            input_transfer = inputTransfer,
        }, JsonOptions);
        try
        {
            AtomicWrite(requestPath, record, overwrite: false);
            return JobOfferResult.Accept(jobId, leaseId, requestDigest);
        }
        catch (IOException)
        {
            return JobOfferResult.Reject("local_job_persist_failed");
        }
    }

    private static JsonElement? ReadTerminal(string path)
    {
        if (!File.Exists(path))
        {
            return null;
        }
        using var document = JsonDocument.Parse(File.ReadAllBytes(path));
        return document.RootElement.Clone();
    }

    private static void AtomicWrite(string path, byte[] content, bool overwrite)
    {
        Directory.CreateDirectory(Path.GetDirectoryName(path)!);
        var temporaryPath = path + ".tmp-" + Guid.NewGuid().ToString("N");
        try
        {
            using (var stream = new FileStream(
                temporaryPath, FileMode.CreateNew, FileAccess.Write, FileShare.None,
                bufferSize: 4096, FileOptions.WriteThrough))
            {
                stream.Write(content);
                stream.Flush(flushToDisk: true);
            }
            File.Move(temporaryPath, path, overwrite);
        }
        finally
        {
            if (File.Exists(temporaryPath)) File.Delete(temporaryPath);
        }
    }

    private static bool IsValidRequest(JsonElement request) =>
        HasOnlyProperties(request, "schema_version", "input", "plot", "output")
        && request.TryGetProperty("schema_version", out var schemaVersion)
        && schemaVersion.TryGetInt32(out var version)
        && version == 1
        && request.TryGetProperty("input", out var input)
        && input.ValueKind == JsonValueKind.Object
        && HasOnlyProperties(input, "input_id", "sheet")
        && input.TryGetProperty("input_id", out var inputId)
        && inputId.ValueKind == JsonValueKind.String
        && !string.IsNullOrWhiteSpace(inputId.GetString())
        && request.TryGetProperty("plot", out var plot)
        && plot.ValueKind == JsonValueKind.Object
        && HasOnlyProperties(
            plot, "type", "x", "y", "template", "title", "x_axis", "y_axis", "legend", "error_bars")
        && plot.TryGetProperty("type", out var plotType)
        && plotType.ValueKind == JsonValueKind.String
        && PlotTypes.Contains(plotType.GetString() ?? "")
        && (!plot.TryGetProperty("title", out var title)
            || title.ValueKind == JsonValueKind.String && title.GetString()!.Length <= 500)
        && plot.TryGetProperty("x", out var x)
        && IsColumnName(x)
        && plot.TryGetProperty("y", out var y)
        && IsValidYColumns(y)
        && (!plot.TryGetProperty("template", out var template)
            || template.GetString() == "publication_double_column")
        && IsValidAxis(plot, "x_axis")
        && IsValidAxis(plot, "y_axis")
        && IsValidLegend(plot)
        && !plot.TryGetProperty("error_bars", out _)
        && request.TryGetProperty("output", out var output)
        && output.ValueKind == JsonValueKind.Object
        && HasOnlyProperties(output, "formats", "dpi", "capture_screenshots", "record_video")
        && output.TryGetProperty("formats", out var formats)
        && formats.ValueKind == JsonValueKind.Array
        && formats.GetArrayLength() > 0
        && IsValidFormats(formats)
        && (!output.TryGetProperty("dpi", out var dpi)
            || dpi.TryGetInt32(out var dpiValue) && dpiValue is >= 72 and <= 1200)
        && IsOptionalBoolean(output, "capture_screenshots")
        && IsOptionalBoolean(output, "record_video")
        && (!output.TryGetProperty("record_video", out var recordVideo)
            || recordVideo.ValueKind == JsonValueKind.False);

    private static bool IsColumnName(JsonElement value) =>
        value.ValueKind == JsonValueKind.String
        && value.GetString() is { Length: >= 1 and <= 128 };

    private static bool IsValidFormats(JsonElement formats)
    {
        var values = formats.EnumerateArray().ToArray();
        return values.All(item =>
                item.ValueKind == JsonValueKind.String
                && OutputFormats.Contains(item.GetString() ?? ""))
            && values.Select(item => item.GetString()).Distinct(StringComparer.Ordinal).Count()
                == values.Length;
    }

    private static bool IsValidYColumns(JsonElement value)
    {
        if (IsColumnName(value)) return true;
        if (value.ValueKind != JsonValueKind.Array
            || value.GetArrayLength() is < 1 or > 16)
        {
            return false;
        }
        var names = value.EnumerateArray().Select(item => item.GetString()).ToArray();
        return value.EnumerateArray().All(IsColumnName)
            && names.Distinct(StringComparer.Ordinal).Count() == names.Length;
    }

    private static bool IsValidAxis(JsonElement plot, string name)
    {
        if (!plot.TryGetProperty(name, out var axis)) return true;
        return axis.ValueKind == JsonValueKind.Object
            && HasOnlyProperties(axis, "title", "unit", "scale")
            && (!axis.TryGetProperty("title", out var title) || title.ValueKind == JsonValueKind.String)
            && (!axis.TryGetProperty("unit", out var unit) || unit.ValueKind == JsonValueKind.String)
            && (!axis.TryGetProperty("scale", out var scale) || scale.GetString() == "linear");
    }

    private static bool IsValidLegend(JsonElement plot)
    {
        if (!plot.TryGetProperty("legend", out var legend)) return true;
        return legend.ValueKind == JsonValueKind.Object
            && HasOnlyProperties(legend, "enabled", "position")
            && (!legend.TryGetProperty("enabled", out var enabled)
                || enabled.ValueKind == JsonValueKind.True)
            && (!legend.TryGetProperty("position", out var position)
                || position.GetString() == "top_right");
    }

    private static bool IsOptionalBoolean(JsonElement value, string name) =>
        !value.TryGetProperty(name, out var property)
        || property.ValueKind is JsonValueKind.True or JsonValueKind.False;

    private static bool IsValidInputTransfer(JsonElement transfer) =>
        transfer.ValueKind == JsonValueKind.Object
        && HasOnlyProperties(transfer, "artifact_id", "filename", "size_bytes", "sha256", "download_path")
        && transfer.TryGetProperty("artifact_id", out var artifactId)
        && Guid.TryParse(artifactId.GetString(), out _)
        && transfer.TryGetProperty("filename", out var filename)
        && filename.ValueKind == JsonValueKind.String
        && Path.GetFileName(filename.GetString()) == filename.GetString()
        && transfer.TryGetProperty("size_bytes", out var size)
        && size.TryGetInt64(out var sizeBytes)
        && sizeBytes is >= 0 and <= 104_857_600
        && transfer.TryGetProperty("sha256", out var sha)
        && sha.GetString() is { Length: 64 }
        && transfer.TryGetProperty("download_path", out var downloadPath)
        && downloadPath.GetString()?.StartsWith("/v1/compute/jobs/", StringComparison.Ordinal) == true
        && downloadPath.GetString()?.EndsWith("/input", StringComparison.Ordinal) == true;

    internal string InputPath(RecoverableJob job)
    {
        if (job.InputTransfer is not JsonElement transfer
            || !IsValidInputTransfer(transfer))
        {
            throw new InvalidDataException("Stored input transfer is invalid.");
        }
        var filename = transfer.GetProperty("filename").GetString()!;
        return Path.Combine(jobsDirectory, job.JobId.ToString("D"), "input", filename);
    }

    internal void WriteTerminal(
        RecoverableJob job, string status, string code, string detail)
    {
        var path = Path.Combine(jobsDirectory, job.JobId.ToString("D"), "terminal.json");
        if (File.Exists(path))
        {
            return;
        }
        var content = JsonSerializer.SerializeToUtf8Bytes(new
        {
            job_id = job.JobId,
            lease_id = job.LeaseId,
            request_digest = job.RequestDigest,
            status,
            error = new { code, detail },
            artifact_manifest = Array.Empty<object>(),
            terminal_at = DateTimeOffset.UtcNow,
        }, JsonOptions);
        AtomicWrite(path, content, overwrite: false);
    }

    private static bool HasOnlyProperties(JsonElement value, params string[] allowed)
    {
        var names = new HashSet<string>(allowed, StringComparer.Ordinal);
        return value.EnumerateObject().All(item => names.Contains(item.Name));
    }

    private static bool TryReadGuid(JsonElement payload, string name, out Guid value)
    {
        value = Guid.Empty;
        return payload.TryGetProperty(name, out var property)
            && property.ValueKind == JsonValueKind.String
            && Guid.TryParse(property.GetString(), out value);
    }
}

internal sealed record RecoverableJob(
    Guid JobId,
    Guid LeaseId,
    string RequestDigest,
    JsonElement? InputTransfer,
    JsonElement? Terminal,
    bool UploadComplete);

internal sealed record JobOfferResult(
    bool Accepted, Guid JobId, Guid LeaseId, string RequestDigest, string Reason)
{
    internal static JobOfferResult Accept(Guid jobId, Guid leaseId, string requestDigest) =>
        new(true, jobId, leaseId, requestDigest, "");

    internal static JobOfferResult Reject(string reason) =>
        new(false, Guid.Empty, Guid.Empty, "", reason);
}