using StellaOps.Concelier.SbomIntegration.Models;
using StellaOps.Scanner.AiMlSecurity.Models;
using StellaOps.Scanner.AiMlSecurity.Policy;
using System.Collections.Immutable;

namespace StellaOps.Scanner.AiMlSecurity.Analyzers;

public sealed class ModelProvenanceVerifier : IAiMlSecurityCheck
{
    public Task<AiMlSecurityResult> AnalyzeAsync(AiMlSecurityContext context, CancellationToken ct = default)
    {
        var findings = new List<AiSecurityFinding>();
        var provenancePolicy = context.Policy.ProvenanceRequirements;

        foreach (var component in context.ModelComponents)
        {
            ct.ThrowIfCancellationRequested();

            if (context.IsExempted(component))
            {
                continue;
            }

            var hasHash = !component.Hashes.IsDefaultOrEmpty;
            var hasSignature = HasSignature(component);
            var source = ResolveSource(component);
            var hasTrustedSource = HasTrustedSource(source, provenancePolicy);

            if ((provenancePolicy.RequireHash && !hasHash)
                || (provenancePolicy.RequireSignature && !hasSignature)
                || (!provenancePolicy.TrustedSources.IsDefaultOrEmpty && !hasTrustedSource))
            {
                findings.Add(new AiSecurityFinding
                {
                    Type = AiSecurityFindingType.UnverifiedModelProvenance,
                    Severity = provenancePolicy.RequireSignature ? Severity.High : Severity.Medium,
                    Title = "Unverified model provenance",
                    Description = "Model provenance does not meet policy requirements.",
                    Remediation = "Provide hashes/signatures and trusted source references.",
                    ComponentName = component.Name,
                    ComponentBomRef = component.BomRef,
                    ModelName = component.Name,
                    Metadata = ImmutableDictionary<string, string>.Empty
                        .Add("hasHash", hasHash.ToString())
                        .Add("hasSignature", hasSignature.ToString())
                        .Add("source", source ?? string.Empty)
                });
            }

            if (component.Modified || HasLineage(component))
            {
                findings.Add(new AiSecurityFinding
                {
                    Type = AiSecurityFindingType.ModelDriftRisk,
                    Severity = Severity.Medium,
                    Title = "Model drift risk",
                    Description = "Model indicates modifications or fine-tuning lineage.",
                    Remediation = "Review fine-tuning lineage and validate drift monitoring.",
                    ComponentName = component.Name,
                    ComponentBomRef = component.BomRef,
                    ModelName = component.Name
                });
            }

            if (IsAdversarialVulnerable(component))
            {
                findings.Add(new AiSecurityFinding
                {
                    Type = AiSecurityFindingType.AdversarialVulnerability,
                    Severity = Severity.High,
                    Title = "Adversarial vulnerability flagged",
                    Description = "Model indicates adversarial robustness concerns.",
                    Remediation = "Perform adversarial testing and mitigation.",
                    ComponentName = component.Name,
                    ComponentBomRef = component.BomRef,
                    ModelName = component.Name
                });
            }
        }

        return Task.FromResult(new AiMlSecurityResult
        {
            Findings = findings.ToImmutableArray()
        });
    }

    private static bool HasSignature(ParsedComponent component)
    {
        if (component.ExternalReferences.Any(reference =>
            (reference.Type ?? string.Empty).Contains("signature", StringComparison.OrdinalIgnoreCase)))
        {
            return true;
        }

        foreach (var pair in component.Properties)
        {
            if (pair.Key.Contains("signature", StringComparison.OrdinalIgnoreCase)
                && IsTruthy(pair.Value))
            {
                return true;
            }
        }

        return false;
    }

    private static string? ResolveSource(ParsedComponent component)
    {
        if (!string.IsNullOrWhiteSpace(component.Publisher))
        {
            return component.Publisher;
        }

        if (!string.IsNullOrWhiteSpace(component.Supplier?.Name))
        {
            return component.Supplier?.Name;
        }

        var external = component.ExternalReferences
            .Select(reference => reference.Url)
            .FirstOrDefault(url => !string.IsNullOrWhiteSpace(url));

        return external;
    }

    private static bool HasTrustedSource(string? source, AiProvenanceRequirements policy)
    {
        if (string.IsNullOrWhiteSpace(source))
        {
            return false;
        }

        var normalized = source.ToLowerInvariant();
        return policy.TrustedSources.Any(entry => normalized.Contains(entry, StringComparison.OrdinalIgnoreCase))
            || policy.KnownModelHubs.Any(entry => normalized.Contains(entry, StringComparison.OrdinalIgnoreCase));
    }

    private static bool HasLineage(ParsedComponent component)
    {
        var pedigree = component.Pedigree;
        if (pedigree is null)
        {
            return false;
        }

        return !pedigree.Ancestors.IsDefaultOrEmpty || !pedigree.Variants.IsDefaultOrEmpty;
    }

    private static bool IsAdversarialVulnerable(ParsedComponent component)
    {
        if (component.Properties.TryGetValue("ai:adversarialVulnerability", out var value)
            && IsTruthy(value))
        {
            return true;
        }

        if (component.Properties.TryGetValue("ai:adversarial", out var shorthand)
            && IsTruthy(shorthand))
        {
            return true;
        }

        if (component.ModelCard?.Considerations?.TechnicalLimitations is { } limitations)
        {
            foreach (var limitation in limitations)
            {
                if (limitation.Contains("adversarial", StringComparison.OrdinalIgnoreCase))
                {
                    return true;
                }
            }
        }

        return false;
    }

    private static bool IsTruthy(string? value)
    {
        if (string.IsNullOrWhiteSpace(value))
        {
            return false;
        }

        return value.Equals("true", StringComparison.OrdinalIgnoreCase)
            || value.Equals("yes", StringComparison.OrdinalIgnoreCase)
            || value.Equals("1", StringComparison.OrdinalIgnoreCase);
    }
}