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

public sealed class CycloneDxNormalizer : IVexNormalizer
{
    private static readonly ImmutableDictionary<string, VexClaimStatus> StateMap = new Dictionary<string, VexClaimStatus>(StringComparer.OrdinalIgnoreCase)
    {
        ["not_affected"] = VexClaimStatus.NotAffected,
        ["resolved"] = VexClaimStatus.Fixed,
        ["resolved_with_patches"] = VexClaimStatus.Fixed,
        ["resolved_no_fix"] = VexClaimStatus.Fixed,
        ["fixed"] = VexClaimStatus.Fixed,
        ["affected"] = VexClaimStatus.Affected,
        ["known_affected"] = VexClaimStatus.Affected,
        ["exploitable"] = VexClaimStatus.Affected,
        ["in_triage"] = VexClaimStatus.UnderInvestigation,
        ["under_investigation"] = VexClaimStatus.UnderInvestigation,
        ["unknown"] = VexClaimStatus.UnderInvestigation,
    }.ToImmutableDictionary(StringComparer.OrdinalIgnoreCase);

    private static readonly ImmutableDictionary<string, VexJustification> JustificationMap = new Dictionary<string, VexJustification>(StringComparer.OrdinalIgnoreCase)
    {
        ["code_not_present"] = VexJustification.CodeNotPresent,
        ["code_not_reachable"] = VexJustification.CodeNotReachable,
        ["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,
        ["requires_configuration"] = VexJustification.RequiresConfiguration,
        ["requires_dependency"] = VexJustification.RequiresDependency,
        ["requires_environment"] = VexJustification.RequiresEnvironment,
    }.ToImmutableDictionary(StringComparer.OrdinalIgnoreCase);

    private readonly ILogger<CycloneDxNormalizer> _logger;

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

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

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

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

        try
        {
            var parseResult = CycloneDxParser.Parse(document);
            var baseMetadata = parseResult.Metadata;
            var claimsBuilder = ImmutableArray.CreateBuilder<VexClaim>();

            foreach (var vulnerability in parseResult.Vulnerabilities)
            {
                cancellationToken.ThrowIfCancellationRequested();

                var state = MapState(vulnerability.AnalysisState, out var stateRaw);
                var justification = MapJustification(vulnerability.AnalysisJustification);
                var responses = vulnerability.AnalysisResponses;

                foreach (var affect in vulnerability.Affects)
                {
                    var productInfo = parseResult.ResolveProduct(affect.ComponentRef);
                    var product = new VexProduct(
                        productInfo.Key,
                        productInfo.Name,
                        productInfo.Version,
                        productInfo.Purl,
                        productInfo.Cpe);

                    var metadata = baseMetadata;
                    if (!string.IsNullOrWhiteSpace(stateRaw))
                    {
                        metadata = metadata.SetItem("cyclonedx.analysis.state", stateRaw);
                    }

                    if (!string.IsNullOrWhiteSpace(vulnerability.AnalysisJustification))
                    {
                        metadata = metadata.SetItem("cyclonedx.analysis.justification", vulnerability.AnalysisJustification);
                    }

                    if (responses.Length > 0)
                    {
                        metadata = metadata.SetItem("cyclonedx.analysis.response", string.Join(",", responses));
                    }

                    if (!string.IsNullOrWhiteSpace(affect.ComponentRef))
                    {
                        metadata = metadata.SetItem("cyclonedx.affects.ref", affect.ComponentRef);
                    }

                    var claimDocument = new VexClaimDocument(
                        VexDocumentFormat.CycloneDx,
                        document.Digest,
                        document.SourceUri,
                        parseResult.BomVersion,
                        signature: null);

                    var claim = new VexClaim(
                        vulnerability.VulnerabilityId,
                        provider.Id,
                        product,
                        state,
                        claimDocument,
                        parseResult.FirstObserved,
                        parseResult.LastObserved,
                        justification,
                        vulnerability.Detail,
                        confidence: null,
                        additionalMetadata: metadata);

                    claimsBuilder.Add(claim);
                }
            }

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

            _logger.LogInformation(
                "Normalized CycloneDX 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 CycloneDX VEX document {SourceUri}", document.SourceUri);
            throw;
        }
    }

    private static VexClaimStatus MapState(string? state, out string? raw)
    {
        raw = state?.Trim();

        if (!string.IsNullOrWhiteSpace(state) && StateMap.TryGetValue(state.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 sealed class CycloneDxParser
    {
        public static CycloneDxParseResult Parse(VexRawDocument document)
        {
            using var json = JsonDocument.Parse(document.Content.ToArray());
            var root = json.RootElement;

            var specVersion = TryGetString(root, "specVersion");
            var bomVersion = TryGetString(root, "version");
            var serialNumber = TryGetString(root, "serialNumber");

            var metadataTimestamp = ParseDate(TryGetProperty(root, "metadata"), "timestamp");
            var observedTimestamp = metadataTimestamp ?? document.RetrievedAt;

            var metadataBuilder = ImmutableDictionary.CreateBuilder<string, string>(StringComparer.Ordinal);
            if (!string.IsNullOrWhiteSpace(specVersion))
            {
                metadataBuilder["cyclonedx.specVersion"] = specVersion!;
            }

            if (!string.IsNullOrWhiteSpace(bomVersion))
            {
                metadataBuilder["cyclonedx.version"] = bomVersion!;
            }

            if (!string.IsNullOrWhiteSpace(serialNumber))
            {
                metadataBuilder["cyclonedx.serialNumber"] = serialNumber!;
            }

            var components = CollectComponents(root);
            var vulnerabilities = CollectVulnerabilities(root);

            return new CycloneDxParseResult(
                metadataBuilder.ToImmutable(),
                bomVersion,
                observedTimestamp,
                observedTimestamp,
                components,
                vulnerabilities);
        }

        private static ImmutableDictionary<string, CycloneDxComponent> CollectComponents(JsonElement root)
        {
            var builder = ImmutableDictionary.CreateBuilder<string, CycloneDxComponent>(StringComparer.Ordinal);

            if (root.TryGetProperty("components", out var components) && components.ValueKind == JsonValueKind.Array)
            {
                foreach (var component in components.EnumerateArray())
                {
                    if (component.ValueKind != JsonValueKind.Object)
                    {
                        continue;
                    }

                    var reference = TryGetString(component, "bom-ref") ?? TryGetString(component, "bomRef");
                    if (string.IsNullOrWhiteSpace(reference))
                    {
                        continue;
                    }

                    var name = TryGetString(component, "name") ?? reference;
                    var version = TryGetString(component, "version");
                    var purl = TryGetString(component, "purl");

                    string? cpe = null;
                    if (component.TryGetProperty("externalReferences", out var externalRefs) && externalRefs.ValueKind == JsonValueKind.Array)
                    {
                        foreach (var referenceEntry in externalRefs.EnumerateArray())
                        {
                            if (referenceEntry.ValueKind != JsonValueKind.Object)
                            {
                                continue;
                            }

                            var type = TryGetString(referenceEntry, "type");
                            if (!string.Equals(type, "cpe", StringComparison.OrdinalIgnoreCase))
                            {
                                continue;
                            }

                            if (referenceEntry.TryGetProperty("url", out var url) && url.ValueKind == JsonValueKind.String)
                            {
                                cpe = url.GetString();
                                break;
                            }
                        }
                    }

                    builder[reference!] = new CycloneDxComponent(reference!, name ?? reference!, version, purl, cpe);
                }
            }

            return builder.ToImmutable();
        }

        private static ImmutableArray<CycloneDxVulnerability> CollectVulnerabilities(JsonElement root)
        {
            if (!root.TryGetProperty("vulnerabilities", out var vulnerabilitiesElement) ||
                vulnerabilitiesElement.ValueKind != JsonValueKind.Array)
            {
                return ImmutableArray<CycloneDxVulnerability>.Empty;
            }

            var builder = ImmutableArray.CreateBuilder<CycloneDxVulnerability>();

            foreach (var vulnerability in vulnerabilitiesElement.EnumerateArray())
            {
                if (vulnerability.ValueKind != JsonValueKind.Object)
                {
                    continue;
                }

                var vulnerabilityId =
                    TryGetString(vulnerability, "id") ??
                    TryGetString(vulnerability, "bom-ref") ??
                    TryGetString(vulnerability, "bomRef") ??
                    TryGetString(vulnerability, "cve");

                if (string.IsNullOrWhiteSpace(vulnerabilityId))
                {
                    continue;
                }

                var detail = TryGetString(vulnerability, "detail") ?? TryGetString(vulnerability, "description");

                var analysis = TryGetProperty(vulnerability, "analysis");
                var analysisState = TryGetString(analysis, "state");
                var analysisJustification = TryGetString(analysis, "justification");
                var analysisResponses = CollectResponses(analysis);

                var affects = CollectAffects(vulnerability);
                if (affects.Length == 0)
                {
                    continue;
                }

                builder.Add(new CycloneDxVulnerability(
                    vulnerabilityId.Trim(),
                    detail?.Trim(),
                    analysisState,
                    analysisJustification,
                    analysisResponses,
                    affects));
            }

            return builder.ToImmutable();
        }

        private static ImmutableArray<string> CollectResponses(JsonElement analysis)
        {
            if (analysis.ValueKind != JsonValueKind.Object ||
                !analysis.TryGetProperty("response", out var responseElement) ||
                responseElement.ValueKind != JsonValueKind.Array)
            {
                return ImmutableArray<string>.Empty;
            }

            var responses = new SortedSet<string>(StringComparer.OrdinalIgnoreCase);
            foreach (var response in responseElement.EnumerateArray())
            {
                if (response.ValueKind == JsonValueKind.String)
                {
                    var value = response.GetString();
                    if (!string.IsNullOrWhiteSpace(value))
                    {
                        responses.Add(value.Trim());
                    }
                }
            }

            return responses.Count == 0 ? ImmutableArray<string>.Empty : responses.ToImmutableArray();
        }

        private static ImmutableArray<CycloneDxAffect> CollectAffects(JsonElement vulnerability)
        {
            if (!vulnerability.TryGetProperty("affects", out var affectsElement) ||
                affectsElement.ValueKind != JsonValueKind.Array)
            {
                return ImmutableArray<CycloneDxAffect>.Empty;
            }

            var builder = ImmutableArray.CreateBuilder<CycloneDxAffect>();
            foreach (var affect in affectsElement.EnumerateArray())
            {
                if (affect.ValueKind != JsonValueKind.Object)
                {
                    continue;
                }

                var reference = TryGetString(affect, "ref");
                if (string.IsNullOrWhiteSpace(reference))
                {
                    continue;
                }

                builder.Add(new CycloneDxAffect(reference.Trim()));
            }

            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 CycloneDxParseResult(
        ImmutableDictionary<string, string> Metadata,
        string? BomVersion,
        DateTimeOffset FirstObserved,
        DateTimeOffset LastObserved,
        ImmutableDictionary<string, CycloneDxComponent> Components,
        ImmutableArray<CycloneDxVulnerability> Vulnerabilities)
    {
        public CycloneDxProductInfo ResolveProduct(string? componentRef)
        {
            if (!string.IsNullOrWhiteSpace(componentRef) &&
                Components.TryGetValue(componentRef.Trim(), out var component))
            {
                return new CycloneDxProductInfo(component.Reference, component.Name, component.Version, component.Purl, component.Cpe);
            }

            var key = string.IsNullOrWhiteSpace(componentRef) ? "unknown-component" : componentRef.Trim();
            return new CycloneDxProductInfo(key, key, null, null, null);
        }
    }

    private sealed record CycloneDxComponent(
        string Reference,
        string Name,
        string? Version,
        string? Purl,
        string? Cpe);

    private sealed record CycloneDxVulnerability(
        string VulnerabilityId,
        string? Detail,
        string? AnalysisState,
        string? AnalysisJustification,
        ImmutableArray<string> AnalysisResponses,
        ImmutableArray<CycloneDxAffect> Affects);

    private sealed record CycloneDxAffect(string ComponentRef);

    private sealed record CycloneDxProductInfo(
        string Key,
        string Name,
        string? Version,
        string? Purl,
        string? Cpe);
}
