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.OpenVEX;

public sealed class OpenVexNormalizer : IVexNormalizer
{
    private static readonly ImmutableDictionary<string, VexClaimStatus> StatusMap = new Dictionary<string, VexClaimStatus>(StringComparer.OrdinalIgnoreCase)
    {
        ["affected"] = VexClaimStatus.Affected,
        ["not_affected"] = VexClaimStatus.NotAffected,
        ["fixed"] = VexClaimStatus.Fixed,
        ["under_investigation"] = VexClaimStatus.UnderInvestigation,
    }.ToImmutableDictionary(StringComparer.OrdinalIgnoreCase);

    private static readonly ImmutableDictionary<string, VexJustification> JustificationMap = new Dictionary<string, VexJustification>(StringComparer.OrdinalIgnoreCase)
    {
        ["component_not_present"] = VexJustification.ComponentNotPresent,
        ["component_not_configured"] = VexJustification.ComponentNotConfigured,
        ["vulnerable_code_not_present"] = VexJustification.VulnerableCodeNotPresent,
        ["vulnerable_code_not_in_execute_path"] = VexJustification.VulnerableCodeNotInExecutePath,
        ["vulnerable_code_cannot_be_controlled_by_adversary"] = VexJustification.VulnerableCodeCannotBeControlledByAdversary,
        ["inline_mitigations_already_exist"] = VexJustification.InlineMitigationsAlreadyExist,
        ["protected_by_mitigating_control"] = VexJustification.ProtectedByMitigatingControl,
        ["protected_by_compensating_control"] = VexJustification.ProtectedByCompensatingControl,
        ["code_not_present"] = VexJustification.CodeNotPresent,
        ["code_not_reachable"] = VexJustification.CodeNotReachable,
        ["requires_configuration"] = VexJustification.RequiresConfiguration,
        ["requires_dependency"] = VexJustification.RequiresDependency,
        ["requires_environment"] = VexJustification.RequiresEnvironment,
    }.ToImmutableDictionary(StringComparer.OrdinalIgnoreCase);

    private readonly ILogger<OpenVexNormalizer> _logger;

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

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

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

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

        try
        {
            var result = OpenVexParser.Parse(document);
            var claims = ImmutableArray.CreateBuilder<VexClaim>(result.Statements.Length);

            foreach (var statement in result.Statements)
            {
                cancellationToken.ThrowIfCancellationRequested();

                var status = MapStatus(statement.Status);
                var justification = MapJustification(statement.Justification);

                foreach (var product in statement.Products)
                {
                    var vexProduct = new VexProduct(
                        product.Key,
                        product.Name,
                        product.Version,
                        product.Purl,
                        product.Cpe);

                    var metadata = result.Metadata;

                    metadata = metadata.SetItem("openvex.statement.id", statement.Id);
                    if (!string.IsNullOrWhiteSpace(statement.Status))
                    {
                        metadata = metadata.SetItem("openvex.statement.status", statement.Status!);
                    }

                    if (!string.IsNullOrWhiteSpace(statement.Justification))
                    {
                        metadata = metadata.SetItem("openvex.statement.justification", statement.Justification!);
                    }

                    if (!string.IsNullOrWhiteSpace(product.OriginalId))
                    {
                        metadata = metadata.SetItem("openvex.product.source", product.OriginalId!);
                    }

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

                    var claim = new VexClaim(
                        statement.Vulnerability,
                        provider.Id,
                        vexProduct,
                        status,
                        claimDocument,
                        result.FirstObserved,
                        result.LastObserved,
                        justification,
                        statement.Remarks,
                        confidence: null,
                        additionalMetadata: 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 OpenVEX 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 OpenVEX document {SourceUri}", document.SourceUri);
            throw;
        }
    }

    private static VexClaimStatus MapStatus(string? status)
    {
        if (!string.IsNullOrWhiteSpace(status) && StatusMap.TryGetValue(status.Trim(), out var mapped))
        {
            return mapped;
        }

        return VexClaimStatus.UnderInvestigation;
    }

    private static VexJustification? MapJustification(string? justification)
    {
        if (string.IsNullOrWhiteSpace(justification))
        {
            return null;
        }

        return JustificationMap.TryGetValue(justification.Trim(), out var mapped)
            ? mapped
            : null;
    }

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

            var metadata = ImmutableDictionary.CreateBuilder<string, string>(StringComparer.Ordinal);

            var documentElement = TryGetProperty(root, "document");
            var version = TryGetString(documentElement, "version");
            var author = TryGetString(documentElement, "author");

            if (!string.IsNullOrWhiteSpace(version))
            {
                metadata["openvex.document.version"] = version!;
            }

            if (!string.IsNullOrWhiteSpace(author))
            {
                metadata["openvex.document.author"] = author!;
            }

            var issued = ParseDate(documentElement, "issued");
            var lastUpdated = ParseDate(documentElement, "last_updated") ?? issued ?? document.RetrievedAt;
            var effectiveDate = ParseDate(documentElement, "effective_date") ?? issued ?? document.RetrievedAt;

            var statements = CollectStatements(root);

            return new OpenVexParseResult(
                metadata.ToImmutable(),
                version,
                effectiveDate,
                lastUpdated,
                statements);
        }

        private static ImmutableArray<OpenVexStatement> CollectStatements(JsonElement root)
        {
            if (!root.TryGetProperty("statements", out var statementsElement) ||
                statementsElement.ValueKind != JsonValueKind.Array)
            {
                return ImmutableArray<OpenVexStatement>.Empty;
            }

            var builder = ImmutableArray.CreateBuilder<OpenVexStatement>();
            foreach (var statement in statementsElement.EnumerateArray())
            {
                if (statement.ValueKind != JsonValueKind.Object)
                {
                    continue;
                }

                var vulnerability = TryGetString(statement, "vulnerability") ?? TryGetString(statement, "vuln") ?? string.Empty;
                if (string.IsNullOrWhiteSpace(vulnerability))
                {
                    continue;
                }

                var id = TryGetString(statement, "id") ?? Guid.NewGuid().ToString();
                var status = TryGetString(statement, "status");
                var justification = TryGetString(statement, "justification");
                var remarks = TryGetString(statement, "remediation") ?? TryGetString(statement, "statement");
                var products = CollectProducts(statement);
                if (products.Length == 0)
                {
                    continue;
                }

                builder.Add(new OpenVexStatement(
                    id,
                    vulnerability.Trim(),
                    status,
                    justification,
                    remarks,
                    products));
            }

            return builder.ToImmutable();
        }

        private static ImmutableArray<OpenVexProduct> CollectProducts(JsonElement statement)
        {
            if (!statement.TryGetProperty("products", out var productsElement) ||
                productsElement.ValueKind != JsonValueKind.Array)
            {
                return ImmutableArray<OpenVexProduct>.Empty;
            }

            var builder = ImmutableArray.CreateBuilder<OpenVexProduct>();
            foreach (var product in productsElement.EnumerateArray())
            {
                if (product.ValueKind != JsonValueKind.String && product.ValueKind != JsonValueKind.Object)
                {
                    continue;
                }

                if (product.ValueKind == JsonValueKind.String)
                {
                    var value = product.GetString();
                    if (string.IsNullOrWhiteSpace(value))
                    {
                        continue;
                    }

                    builder.Add(OpenVexProduct.FromString(value.Trim()));
                    continue;
                }

                var id = TryGetString(product, "id") ?? TryGetString(product, "product_id");
                var name = TryGetString(product, "name");
                var version = TryGetString(product, "version");
                var purl = TryGetString(product, "purl");
                var cpe = TryGetString(product, "cpe");

                if (string.IsNullOrWhiteSpace(id) && string.IsNullOrWhiteSpace(purl))
                {
                    continue;
                }

                builder.Add(new OpenVexProduct(
                    id ?? purl!,
                    name ?? id ?? purl!,
                    version,
                    purl,
                    cpe,
                    OriginalId: id));
            }

            return builder.ToImmutable();
        }

        private static JsonElement TryGetProperty(JsonElement element, string propertyName)
            => element.ValueKind == JsonValueKind.Object && element.TryGetProperty(propertyName, out var value)
                ? value
                : default;

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

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

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

        private static DateTimeOffset? ParseDate(JsonElement element, string propertyName)
        {
            var value = TryGetString(element, propertyName);
            if (string.IsNullOrWhiteSpace(value))
            {
                return null;
            }

            return DateTimeOffset.TryParse(value, out var parsed) ? parsed : null;
        }
    }

    private sealed record OpenVexParseResult(
        ImmutableDictionary<string, string> Metadata,
        string? DocumentVersion,
        DateTimeOffset FirstObserved,
        DateTimeOffset LastObserved,
        ImmutableArray<OpenVexStatement> Statements);

    private sealed record OpenVexStatement(
        string Id,
        string Vulnerability,
        string? Status,
        string? Justification,
        string? Remarks,
        ImmutableArray<OpenVexProduct> Products);

    private sealed record OpenVexProduct(
        string Key,
        string Name,
        string? Version,
        string? Purl,
        string? Cpe,
        string? OriginalId)
    {
        public static OpenVexProduct FromString(string value)
        {
            var key = value;
            string? purl = null;
            string? name = value;

            if (value.StartsWith("pkg:", StringComparison.OrdinalIgnoreCase))
            {
                purl = value;
            }

            return new OpenVexProduct(key, name, null, purl, null, OriginalId: value);
        }
    }
}
