using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.Logging;
using StellaOps.Excititor.Core;

namespace StellaOps.Excititor.Formats.CSAF;

public sealed class CsafNormalizer : IVexNormalizer
{
    private static readonly ImmutableDictionary<VexClaimStatus, int> StatusPrecedence = new Dictionary<VexClaimStatus, int>
    {
        [VexClaimStatus.UnderInvestigation] = 0,
        [VexClaimStatus.Affected] = 1,
        [VexClaimStatus.NotAffected] = 2,
        [VexClaimStatus.Fixed] = 3,
    }.ToImmutableDictionary();

    private readonly ILogger<CsafNormalizer> _logger;

    public CsafNormalizer(ILogger<CsafNormalizer> logger)
    {
        _logger = logger ?? throw new ArgumentNullException(nameof(logger));
    }

    public string Format => VexDocumentFormat.Csaf.ToString().ToLowerInvariant();

    public bool CanHandle(VexRawDocument document)
        => document is not null && document.Format == VexDocumentFormat.Csaf;

    public ValueTask<VexClaimBatch> NormalizeAsync(VexRawDocument document, VexProvider provider, CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(document);
        ArgumentNullException.ThrowIfNull(provider);

        cancellationToken.ThrowIfCancellationRequested();

        try
        {
            var result = CsafParser.Parse(document);
            var claims = ImmutableArray.CreateBuilder<VexClaim>(result.Claims.Length);
            foreach (var entry in result.Claims)
            {
                var product = new VexProduct(
                    entry.Product.ProductId,
                    entry.Product.Name,
                    entry.Product.Version,
                    entry.Product.Purl,
                    entry.Product.Cpe);

                var claimDocument = new VexClaimDocument(
                    VexDocumentFormat.Csaf,
                    document.Digest,
                    document.SourceUri,
                    result.Revision,
                    signature: null);

                var claim = new VexClaim(
                    entry.VulnerabilityId,
                    provider.Id,
                    product,
                    entry.Status,
                    claimDocument,
                    result.FirstRelease,
                    result.LastRelease,
                    justification: null,
                    detail: entry.Detail,
                    confidence: null,
                    additionalMetadata: result.Metadata);

                claims.Add(claim);
            }

            var orderedClaims = claims
                .ToImmutable()
                .OrderBy(static claim => claim.VulnerabilityId, StringComparer.Ordinal)
                .ThenBy(static claim => claim.Product.Key, StringComparer.Ordinal)
                .ToImmutableArray();

            _logger.LogInformation(
                "Normalized CSAF document {Source} into {ClaimCount} claim(s).",
                document.SourceUri,
                orderedClaims.Length);

            return ValueTask.FromResult(new VexClaimBatch(
                document,
                orderedClaims,
                ImmutableDictionary<string, string>.Empty));
        }
        catch (JsonException ex)
        {
            _logger.LogError(ex, "Failed to parse CSAF document {SourceUri}", document.SourceUri);
            throw;
        }
    }

    private static class CsafParser
    {
        public static CsafParseResult Parse(VexRawDocument document)
        {
            using var json = JsonDocument.Parse(document.Content.ToArray());
            var root = json.RootElement;

            var tracking = root.TryGetProperty("document", out var documentElement) &&
                           documentElement.ValueKind == JsonValueKind.Object &&
                           documentElement.TryGetProperty("tracking", out var trackingElement)
                ? trackingElement
                : default;

            var firstRelease = ParseDate(tracking, "initial_release_date") ?? document.RetrievedAt;
            var lastRelease = ParseDate(tracking, "current_release_date") ?? firstRelease;

            if (lastRelease < firstRelease)
            {
                lastRelease = firstRelease;
            }

            var metadataBuilder = ImmutableDictionary.CreateBuilder<string, string>(StringComparer.Ordinal);
            AddIfPresent(metadataBuilder, tracking, "id", "csaf.tracking.id");
            AddIfPresent(metadataBuilder, tracking, "version", "csaf.tracking.version");
            AddIfPresent(metadataBuilder, tracking, "status", "csaf.tracking.status");
            AddPublisherMetadata(metadataBuilder, documentElement);

            var revision = TryGetString(tracking, "revision");

            var productCatalog = CollectProducts(root);
            var claimsBuilder = ImmutableArray.CreateBuilder<CsafClaimEntry>();

            if (root.TryGetProperty("vulnerabilities", out var vulnerabilitiesElement) &&
                vulnerabilitiesElement.ValueKind == JsonValueKind.Array)
            {
                foreach (var vulnerability in vulnerabilitiesElement.EnumerateArray())
                {
                    var vulnerabilityId = ResolveVulnerabilityId(vulnerability);
                    if (string.IsNullOrWhiteSpace(vulnerabilityId))
                    {
                        continue;
                    }

                    var detail = ResolveDetail(vulnerability);
                    var productClaims = BuildClaimsForVulnerability(
                        vulnerabilityId,
                        vulnerability,
                        productCatalog,
                        detail);

                    claimsBuilder.AddRange(productClaims);
                }
            }

            return new CsafParseResult(
                firstRelease,
                lastRelease,
                revision,
                metadataBuilder.ToImmutable(),
                claimsBuilder.ToImmutable());
        }

        private static IReadOnlyList<CsafClaimEntry> BuildClaimsForVulnerability(
            string vulnerabilityId,
            JsonElement vulnerability,
            IReadOnlyDictionary<string, CsafProductInfo> productCatalog,
            string? detail)
        {
            if (!vulnerability.TryGetProperty("product_status", out var statusElement) ||
                statusElement.ValueKind != JsonValueKind.Object)
            {
                return Array.Empty<CsafClaimEntry>();
            }

            var claims = new Dictionary<string, CsafClaimEntryBuilder>(StringComparer.OrdinalIgnoreCase);

            foreach (var statusProperty in statusElement.EnumerateObject())
            {
                var status = MapStatus(statusProperty.Name);
                if (status is null)
                {
                    continue;
                }

                if (statusProperty.Value.ValueKind != JsonValueKind.Array)
                {
                    continue;
                }

                foreach (var productIdElement in statusProperty.Value.EnumerateArray())
                {
                    var productId = productIdElement.GetString();
                    if (string.IsNullOrWhiteSpace(productId))
                    {
                        continue;
                    }

                    var product = ResolveProduct(productCatalog, productId);
                    UpdateClaim(claims, product, status.Value, detail);
                }
            }

            if (claims.Count == 0)
            {
                return Array.Empty<CsafClaimEntry>();
            }

            return claims.Values
                .Select(builder => new CsafClaimEntry(
                    vulnerabilityId,
                    builder.Product,
                    builder.Status,
                    builder.Detail))
                .ToArray();
        }

        private static void UpdateClaim(
            IDictionary<string, CsafClaimEntryBuilder> claims,
            CsafProductInfo product,
            VexClaimStatus status,
            string? detail)
        {
            if (!claims.TryGetValue(product.ProductId, out var existing) ||
                StatusPrecedence[status] > StatusPrecedence[existing.Status])
            {
                claims[product.ProductId] = new CsafClaimEntryBuilder(product, status, detail);
            }
        }

        private static CsafProductInfo ResolveProduct(
            IReadOnlyDictionary<string, CsafProductInfo> catalog,
            string productId)
        {
            if (catalog.TryGetValue(productId, out var product))
            {
                return product;
            }

            return new CsafProductInfo(productId, productId, null, null, null);
        }

        private static string ResolveVulnerabilityId(JsonElement vulnerability)
        {
            var id = TryGetString(vulnerability, "cve")
                ?? TryGetString(vulnerability, "id")
                ?? TryGetString(vulnerability, "vuln_id");

            return string.IsNullOrWhiteSpace(id) ? string.Empty : id.Trim();
        }

        private static string? ResolveDetail(JsonElement vulnerability)
        {
            var title = TryGetString(vulnerability, "title");
            if (!string.IsNullOrWhiteSpace(title))
            {
                return title.Trim();
            }

            if (vulnerability.TryGetProperty("notes", out var notesElement) &&
                notesElement.ValueKind == JsonValueKind.Array)
            {
                foreach (var note in notesElement.EnumerateArray())
                {
                    if (note.ValueKind != JsonValueKind.Object)
                    {
                        continue;
                    }

                    var category = TryGetString(note, "category");
                    if (!string.IsNullOrWhiteSpace(category) &&
                        !string.Equals(category, "description", StringComparison.OrdinalIgnoreCase))
                    {
                        continue;
                    }

                    var text = TryGetString(note, "text");
                    if (!string.IsNullOrWhiteSpace(text))
                    {
                        return text.Trim();
                    }
                }
            }

            return null;
        }

        private static Dictionary<string, CsafProductInfo> CollectProducts(JsonElement root)
        {
            var products = new Dictionary<string, CsafProductInfo>(StringComparer.OrdinalIgnoreCase);

            if (!root.TryGetProperty("product_tree", out var productTree) ||
                productTree.ValueKind != JsonValueKind.Object)
            {
                return products;
            }

            if (productTree.TryGetProperty("full_product_names", out var fullNames) &&
                fullNames.ValueKind == JsonValueKind.Array)
            {
                foreach (var productEntry in fullNames.EnumerateArray())
                {
                    var product = ParseProduct(productEntry, parentBranchName: null);
                    if (product is not null)
                    {
                        AddOrUpdate(product);
                    }
                }
            }

            if (productTree.TryGetProperty("branches", out var branches) &&
                branches.ValueKind == JsonValueKind.Array)
            {
                foreach (var branch in branches.EnumerateArray())
                {
                    VisitBranch(branch, parentBranchName: null);
                }
            }

            return products;

            void VisitBranch(JsonElement branch, string? parentBranchName)
            {
                if (branch.ValueKind != JsonValueKind.Object)
                {
                    return;
                }

                var branchName = TryGetString(branch, "name") ?? parentBranchName;

                if (branch.TryGetProperty("product", out var productElement))
                {
                    var product = ParseProduct(productElement, branchName);
                    if (product is not null)
                    {
                        AddOrUpdate(product);
                    }
                }

                if (branch.TryGetProperty("branches", out var childBranches) &&
                    childBranches.ValueKind == JsonValueKind.Array)
                {
                    foreach (var childBranch in childBranches.EnumerateArray())
                    {
                        VisitBranch(childBranch, branchName);
                    }
                }
            }

            void AddOrUpdate(CsafProductInfo product)
            {
                if (products.TryGetValue(product.ProductId, out var existing))
                {
                    products[product.ProductId] = MergeProducts(existing, product);
                }
                else
                {
                    products[product.ProductId] = product;
                }
            }

            static CsafProductInfo MergeProducts(CsafProductInfo existing, CsafProductInfo incoming)
            {
                static string ChooseName(string incoming, string fallback)
                    => string.IsNullOrWhiteSpace(incoming) ? fallback : incoming;

                static string? ChooseOptional(string? incoming, string? fallback)
                    => string.IsNullOrWhiteSpace(incoming) ? fallback : incoming;

                return new CsafProductInfo(
                    existing.ProductId,
                    ChooseName(incoming.Name, existing.Name),
                    ChooseOptional(incoming.Version, existing.Version),
                    ChooseOptional(incoming.Purl, existing.Purl),
                    ChooseOptional(incoming.Cpe, existing.Cpe));
            }
        }

        private static CsafProductInfo? ParseProduct(JsonElement element, string? parentBranchName)
        {
            if (element.ValueKind != JsonValueKind.Object)
            {
                return null;
            }

            JsonElement productElement = element;
            if (!element.TryGetProperty("product_id", out var idElement) &&
                element.TryGetProperty("product", out var nestedProduct) &&
                nestedProduct.ValueKind == JsonValueKind.Object &&
                nestedProduct.TryGetProperty("product_id", out idElement))
            {
                productElement = nestedProduct;
            }

            var productId = idElement.GetString();
            if (string.IsNullOrWhiteSpace(productId))
            {
                return null;
            }

            var name = TryGetString(productElement, "name")
                ?? TryGetString(element, "name")
                ?? parentBranchName
                ?? productId;

            var version = TryGetString(productElement, "product_version")
                ?? TryGetString(productElement, "version")
                ?? TryGetString(element, "product_version");

            string? cpe = null;
            string? purl = null;
            if (productElement.TryGetProperty("product_identification_helper", out var helper) &&
                helper.ValueKind == JsonValueKind.Object)
            {
                cpe = TryGetString(helper, "cpe");
                purl = TryGetString(helper, "purl");
            }

            return new CsafProductInfo(productId.Trim(), name.Trim(), version?.Trim(), purl?.Trim(), cpe?.Trim());
        }

        private static VexClaimStatus? MapStatus(string statusName)
        {
            if (string.IsNullOrWhiteSpace(statusName))
            {
                return null;
            }

            return statusName switch
            {
                "known_affected" or "fixed_after_release" or "first_affected" or "last_affected" => VexClaimStatus.Affected,
                "known_not_affected" or "last_not_affected" or "first_not_affected" => VexClaimStatus.NotAffected,
                "fixed" or "first_fixed" or "last_fixed" => VexClaimStatus.Fixed,
                "under_investigation" or "investigating" => VexClaimStatus.UnderInvestigation,
                _ => null,
            };
        }

        private static DateTimeOffset? ParseDate(JsonElement element, string propertyName)
        {
            if (element.ValueKind != JsonValueKind.Object)
            {
                return null;
            }

            if (!element.TryGetProperty(propertyName, out var dateElement))
            {
                return null;
            }

            var value = dateElement.GetString();
            if (string.IsNullOrWhiteSpace(value))
            {
                return null;
            }

            if (DateTimeOffset.TryParse(value, out var parsed))
            {
                return parsed;
            }

            return null;
        }

        private static string? TryGetString(JsonElement element, string propertyName)
        {
            if (element.ValueKind != JsonValueKind.Object)
            {
                return null;
            }

            if (!element.TryGetProperty(propertyName, out var property))
            {
                return null;
            }

            return property.ValueKind == JsonValueKind.String ? property.GetString() : null;
        }

        private static void AddIfPresent(
            ImmutableDictionary<string, string>.Builder builder,
            JsonElement element,
            string propertyName,
            string metadataKey)
        {
            var value = TryGetString(element, propertyName);
            if (!string.IsNullOrWhiteSpace(value))
            {
                builder[metadataKey] = value.Trim();
            }
        }

        private static void AddPublisherMetadata(
            ImmutableDictionary<string, string>.Builder builder,
            JsonElement documentElement)
        {
            if (documentElement.ValueKind != JsonValueKind.Object ||
                !documentElement.TryGetProperty("publisher", out var publisher) ||
                publisher.ValueKind != JsonValueKind.Object)
            {
                return;
            }

            AddIfPresent(builder, publisher, "name", "csaf.publisher.name");
            AddIfPresent(builder, publisher, "category", "csaf.publisher.category");
        }

        private readonly record struct CsafClaimEntryBuilder(
            CsafProductInfo Product,
            VexClaimStatus Status,
            string? Detail);
    }

    private sealed record CsafParseResult(
        DateTimeOffset FirstRelease,
        DateTimeOffset LastRelease,
        string? Revision,
        ImmutableDictionary<string, string> Metadata,
        ImmutableArray<CsafClaimEntry> Claims);

    private sealed record CsafClaimEntry(
        string VulnerabilityId,
        CsafProductInfo Product,
        VexClaimStatus Status,
        string? Detail);

    private sealed record CsafProductInfo(
        string ProductId,
        string Name,
        string? Version,
        string? Purl,
        string? Cpe);
}
