using System;
using System.Collections.Immutable;
using System.IO;
using System.Text;
using System.Globalization;
using Microsoft.Extensions.Logging.Abstractions;
using MongoDB.Driver;
using StellaOps.Excititor.Core;
using StellaOps.Excititor.Export;
using StellaOps.Excititor.Policy;
using StellaOps.Excititor.Storage.Mongo;
using Xunit;

namespace StellaOps.Excititor.Export.Tests;

public sealed class ExportEngineTests
{
    [Fact]
    public async Task ExportAsync_GeneratesAndCachesManifest()
    {
        var store = new InMemoryExportStore();
        var evaluator = new StaticPolicyEvaluator("baseline/v1");
        var dataSource = new InMemoryExportDataSource();
        var exporter = new DummyExporter(VexExportFormat.Json);
        var engine = new VexExportEngine(store, evaluator, dataSource, new[] { exporter }, NullLogger<VexExportEngine>.Instance);

        var query = VexQuery.Create(new[] { new VexQueryFilter("vulnId", "CVE-2025-0001") });
        var context = new VexExportRequestContext(query, VexExportFormat.Json, DateTimeOffset.UtcNow, ForceRefresh: false);

        var manifest = await engine.ExportAsync(context, CancellationToken.None);

        Assert.False(manifest.FromCache);
        Assert.Equal(VexExportFormat.Json, manifest.Format);
        Assert.Equal("baseline/v1", manifest.ConsensusRevision);
        Assert.Equal(1, manifest.ClaimCount);
        Assert.NotNull(dataSource.LastDataSet);
        var expectedEnvelopes = VexExportEnvelopeBuilder.Build(
            dataSource.LastDataSet!,
            VexPolicySnapshot.Default,
            context.RequestedAt);
        Assert.NotNull(manifest.ConsensusDigest);
        Assert.Equal(expectedEnvelopes.ConsensusDigest.Algorithm, manifest.ConsensusDigest!.Algorithm);
        Assert.Equal(expectedEnvelopes.ConsensusDigest.Digest, manifest.ConsensusDigest.Digest);
        Assert.NotNull(manifest.ScoreDigest);
        Assert.Equal(expectedEnvelopes.ScoreDigest.Algorithm, manifest.ScoreDigest!.Algorithm);
        Assert.Equal(expectedEnvelopes.ScoreDigest.Digest, manifest.ScoreDigest.Digest);
        Assert.Empty(manifest.QuietProvenance);

        // second call hits cache
        var cached = await engine.ExportAsync(context, CancellationToken.None);
        Assert.True(cached.FromCache);
        Assert.Equal(manifest.ExportId, cached.ExportId);
    }

    [Fact]
    public async Task ExportAsync_ForceRefreshInvalidatesCacheEntry()
    {
        var store = new InMemoryExportStore();
        var evaluator = new StaticPolicyEvaluator("baseline/v1");
        var dataSource = new InMemoryExportDataSource();
        var exporter = new DummyExporter(VexExportFormat.Json);
        var cacheIndex = new RecordingCacheIndex();
        var engine = new VexExportEngine(store, evaluator, dataSource, new[] { exporter }, NullLogger<VexExportEngine>.Instance, cacheIndex);

        var query = VexQuery.Create(new[] { new VexQueryFilter("vulnId", "CVE-2025-0001") });
        var initialContext = new VexExportRequestContext(query, VexExportFormat.Json, DateTimeOffset.UtcNow);
        _ = await engine.ExportAsync(initialContext, CancellationToken.None);

        var refreshContext = new VexExportRequestContext(query, VexExportFormat.Json, DateTimeOffset.UtcNow.AddMinutes(1), ForceRefresh: true);
        var refreshed = await engine.ExportAsync(refreshContext, CancellationToken.None);

        Assert.False(refreshed.FromCache);
        var signature = VexQuerySignature.FromQuery(refreshContext.Query);
        Assert.True(cacheIndex.RemoveCalls.TryGetValue((signature.Value, refreshContext.Format), out var removed));
        Assert.True(removed);
    }

    [Fact]
    public async Task ExportAsync_WritesArtifactsToAllStores()
    {
        var store = new InMemoryExportStore();
        var evaluator = new StaticPolicyEvaluator("baseline/v1");
        var dataSource = new InMemoryExportDataSource();
        var exporter = new DummyExporter(VexExportFormat.Json);
        var recorder1 = new RecordingArtifactStore();
        var recorder2 = new RecordingArtifactStore();
        var engine = new VexExportEngine(
            store,
            evaluator,
            dataSource,
            new[] { exporter },
            NullLogger<VexExportEngine>.Instance,
            cacheIndex: null,
            artifactStores: new[] { recorder1, recorder2 });

        var query = VexQuery.Create(new[] { new VexQueryFilter("vulnId", "CVE-2025-0001") });
        var context = new VexExportRequestContext(query, VexExportFormat.Json, DateTimeOffset.UtcNow);

        await engine.ExportAsync(context, CancellationToken.None);

        Assert.Equal(1, recorder1.SaveCount);
        Assert.Equal(1, recorder2.SaveCount);
    }

    [Fact]
    public async Task ExportAsync_AttachesAttestationMetadata()
    {
        var store = new InMemoryExportStore();
        var evaluator = new StaticPolicyEvaluator("baseline/v1");
        var dataSource = new InMemoryExportDataSource();
        var exporter = new DummyExporter(VexExportFormat.Json);
        var attestation = new RecordingAttestationClient();
        var engine = new VexExportEngine(
            store,
            evaluator,
            dataSource,
            new[] { exporter },
            NullLogger<VexExportEngine>.Instance,
            cacheIndex: null,
            artifactStores: null,
            attestationClient: attestation);

        var query = VexQuery.Create(new[] { new VexQueryFilter("vulnId", "CVE-2025-0001") });
        var requestedAt = DateTimeOffset.UtcNow;
        var context = new VexExportRequestContext(query, VexExportFormat.Json, requestedAt);

        var manifest = await engine.ExportAsync(context, CancellationToken.None);

        Assert.NotNull(attestation.LastRequest);
        Assert.NotNull(dataSource.LastDataSet);
        var expectedEnvelopes = VexExportEnvelopeBuilder.Build(
            dataSource.LastDataSet!,
            VexPolicySnapshot.Default,
            requestedAt);
        Assert.Equal(manifest.ExportId, attestation.LastRequest!.ExportId);
        var metadata = attestation.LastRequest.Metadata;
        Assert.True(metadata.ContainsKey("consensusDigest"), "Consensus digest metadata missing");
        Assert.Equal(expectedEnvelopes.ConsensusDigest.ToUri(), metadata["consensusDigest"]);
        Assert.True(metadata.ContainsKey("scoreDigest"), "Score digest metadata missing");
        Assert.Equal(expectedEnvelopes.ScoreDigest.ToUri(), metadata["scoreDigest"]);
        Assert.Equal(expectedEnvelopes.Consensus.Length.ToString(CultureInfo.InvariantCulture), metadata["consensusEntryCount"]);
        Assert.Equal(expectedEnvelopes.ScoreEnvelope.Entries.Length.ToString(CultureInfo.InvariantCulture), metadata["scoreEntryCount"]);
        Assert.Equal(VexPolicySnapshot.Default.RevisionId, metadata["policyRevisionId"]);
        Assert.Equal(VexPolicySnapshot.Default.Version, metadata["policyVersion"]);
        Assert.Equal(VexPolicySnapshot.Default.ConsensusOptions.Alpha.ToString("G17", CultureInfo.InvariantCulture), metadata["scoreAlpha"]);
        Assert.Equal(VexPolicySnapshot.Default.ConsensusOptions.Beta.ToString("G17", CultureInfo.InvariantCulture), metadata["scoreBeta"]);
        Assert.Equal(VexPolicySnapshot.Default.ConsensusOptions.WeightCeiling.ToString("G17", CultureInfo.InvariantCulture), metadata["scoreWeightCeiling"]);
        Assert.NotNull(manifest.Attestation);
        Assert.Equal(attestation.Response.Attestation.EnvelopeDigest, manifest.Attestation!.EnvelopeDigest);
        Assert.Equal(attestation.Response.Attestation.PredicateType, manifest.Attestation.PredicateType);
        Assert.NotNull(manifest.ConsensusDigest);
        Assert.Equal(expectedEnvelopes.ConsensusDigest.Digest, manifest.ConsensusDigest!.Digest);
        Assert.NotNull(manifest.ScoreDigest);
        Assert.Equal(expectedEnvelopes.ScoreDigest.Digest, manifest.ScoreDigest!.Digest);
        Assert.Empty(manifest.QuietProvenance);

        Assert.NotNull(store.LastSavedManifest);
        Assert.Equal(manifest.Attestation, store.LastSavedManifest!.Attestation);
        Assert.Equal(manifest.QuietProvenance, store.LastSavedManifest!.QuietProvenance);
    }

    [Fact]
    public async Task ExportAsync_IncludesQuietProvenanceMetadata()
    {
        var store = new InMemoryExportStore();
        var evaluator = new StaticPolicyEvaluator("baseline/v1");
        var dataSource = new QuietExportDataSource();
        var exporter = new DummyExporter(VexExportFormat.Json);
        var attestation = new RecordingAttestationClient();
        var engine = new VexExportEngine(
            store,
            evaluator,
            dataSource,
            new[] { exporter },
            NullLogger<VexExportEngine>.Instance,
            cacheIndex: null,
            artifactStores: null,
            attestationClient: attestation);

        var query = VexQuery.Create(new[] { new VexQueryFilter("vulnId", "CVE-2025-0002") });
        var requestedAt = DateTimeOffset.UtcNow;
        var context = new VexExportRequestContext(query, VexExportFormat.Json, requestedAt);

        var manifest = await engine.ExportAsync(context, CancellationToken.None);

        var quiet = Assert.Single(manifest.QuietProvenance);
        Assert.Equal("CVE-2025-0002", quiet.VulnerabilityId);
        Assert.Equal("pkg:demo/app", quiet.ProductKey);
        var statement = Assert.Single(quiet.Statements);
        Assert.Equal("vendor", statement.ProviderId);
        Assert.Equal("sha256:quiet", statement.StatementId);
        Assert.Equal(VexJustification.ComponentNotPresent, statement.Justification);
        Assert.NotNull(statement.Signature);
        Assert.Equal("quiet-signer", statement.Signature!.Subject);
        Assert.Equal("quiet-key", statement.Signature.KeyId);

        var expectedQuietJson = VexCanonicalJsonSerializer.Serialize(manifest.QuietProvenance);
        Assert.NotNull(attestation.LastRequest);
        Assert.True(attestation.LastRequest!.Metadata.TryGetValue("quietedBy", out var quietJson));
        Assert.Equal(expectedQuietJson, quietJson);
        Assert.True(attestation.LastRequest.Metadata.TryGetValue("quietedByStatementCount", out var quietCount));
        Assert.Equal("1", quietCount);

        Assert.NotNull(store.LastSavedManifest);
        Assert.Equal(manifest.QuietProvenance, store.LastSavedManifest!.QuietProvenance);
    }

    private sealed class InMemoryExportStore : IVexExportStore
    {
        private readonly Dictionary<string, VexExportManifest> _store = new(StringComparer.Ordinal);

        public VexExportManifest? LastSavedManifest { get; private set; }

        public ValueTask<VexExportManifest?> FindAsync(VexQuerySignature signature, VexExportFormat format, CancellationToken cancellationToken, IClientSessionHandle? session = null)
        {
            var key = CreateKey(signature.Value, format);
            _store.TryGetValue(key, out var manifest);
            return ValueTask.FromResult<VexExportManifest?>(manifest);
        }

        public ValueTask SaveAsync(VexExportManifest manifest, CancellationToken cancellationToken, IClientSessionHandle? session = null)
        {
            var key = CreateKey(manifest.QuerySignature.Value, manifest.Format);
            _store[key] = manifest;
            LastSavedManifest = manifest;
            return ValueTask.CompletedTask;
        }

        private static string CreateKey(string signature, VexExportFormat format)
            => FormattableString.Invariant($"{signature}|{format}");
    }

    private sealed class QuietExportDataSource : IVexExportDataSource
    {
        public ValueTask<VexExportDataSet> FetchAsync(VexQuery query, CancellationToken cancellationToken)
        {
            var signature = new VexSignatureMetadata(
                type: "pgp",
                subject: "quiet-signer",
                issuer: "quiet-ca",
                keyId: "quiet-key",
                verifiedAt: DateTimeOffset.UnixEpoch,
                transparencyLogReference: "rekor://quiet");

            var claim = new VexClaim(
                "CVE-2025-0002",
                "vendor",
                new VexProduct("pkg:demo/app", "Demo"),
                VexClaimStatus.NotAffected,
                new VexClaimDocument(VexDocumentFormat.OpenVex, "sha256:quiet", new Uri("https://example.org/quiet"), signature: signature),
                DateTimeOffset.UtcNow,
                DateTimeOffset.UtcNow,
                justification: VexJustification.ComponentNotPresent);

            var consensus = new VexConsensus(
                "CVE-2025-0002",
                claim.Product,
                VexConsensusStatus.NotAffected,
                DateTimeOffset.UtcNow,
                new[]
                {
                    new VexConsensusSource("vendor", VexClaimStatus.NotAffected, "sha256:quiet", 1.0, claim.Justification),
                },
                conflicts: null,
                policyVersion: "baseline/v1",
                summary: "not_affected");

            return ValueTask.FromResult(new VexExportDataSet(
                ImmutableArray.Create(consensus),
                ImmutableArray.Create(claim),
                ImmutableArray.Create("vendor")));
        }
    }

    private sealed class RecordingAttestationClient : IVexAttestationClient
    {
        public VexAttestationRequest? LastRequest { get; private set; }

        public VexAttestationResponse Response { get; } = new VexAttestationResponse(
            new VexAttestationMetadata(
                predicateType: "https://stella-ops.org/attestations/vex-export",
                rekor: new VexRekorReference("0.2", "rekor://entry", "123"),
                envelopeDigest: "sha256:envelope",
                signedAt: DateTimeOffset.UnixEpoch),
            ImmutableDictionary<string, string>.Empty);

        public ValueTask<VexAttestationResponse> SignAsync(VexAttestationRequest request, CancellationToken cancellationToken)
        {
            LastRequest = request;
            return ValueTask.FromResult(Response);
        }

        public ValueTask<VexAttestationVerification> VerifyAsync(VexAttestationRequest request, CancellationToken cancellationToken)
            => ValueTask.FromResult(new VexAttestationVerification(true, ImmutableDictionary<string, string>.Empty));
    }

    private sealed class RecordingCacheIndex : IVexCacheIndex
    {
        public Dictionary<(string Signature, VexExportFormat Format), bool> RemoveCalls { get; } = new();

        public ValueTask<VexCacheEntry?> FindAsync(VexQuerySignature signature, VexExportFormat format, CancellationToken cancellationToken, IClientSessionHandle? session = null)
            => ValueTask.FromResult<VexCacheEntry?>(null);

        public ValueTask SaveAsync(VexCacheEntry entry, CancellationToken cancellationToken, IClientSessionHandle? session = null)
            => ValueTask.CompletedTask;

        public ValueTask RemoveAsync(VexQuerySignature signature, VexExportFormat format, CancellationToken cancellationToken, IClientSessionHandle? session = null)
        {
            RemoveCalls[(signature.Value, format)] = true;
            return ValueTask.CompletedTask;
        }
    }

    private sealed class RecordingArtifactStore : IVexArtifactStore
    {
        public int SaveCount { get; private set; }

        public ValueTask<VexStoredArtifact> SaveAsync(VexExportArtifact artifact, CancellationToken cancellationToken)
        {
            SaveCount++;
            return ValueTask.FromResult(new VexStoredArtifact(artifact.ContentAddress, "memory", artifact.Content.Length, artifact.Metadata));
        }

        public ValueTask DeleteAsync(VexContentAddress contentAddress, CancellationToken cancellationToken)
            => ValueTask.CompletedTask;

        public ValueTask<Stream?> OpenReadAsync(VexContentAddress contentAddress, CancellationToken cancellationToken)
            => ValueTask.FromResult<Stream?>(null);
    }

    private sealed class StaticPolicyEvaluator : IVexPolicyEvaluator
    {
        public StaticPolicyEvaluator(string version)
        {
            Version = version;
        }

        public string Version { get; }

        public VexPolicySnapshot Snapshot => VexPolicySnapshot.Default;

        public double GetProviderWeight(VexProvider provider) => 1.0;

        public bool IsClaimEligible(VexClaim claim, VexProvider provider, out string? rejectionReason)
        {
            rejectionReason = null;
            return true;
        }
    }

    private sealed class InMemoryExportDataSource : IVexExportDataSource
    {
        public VexExportDataSet? LastDataSet { get; private set; }

        public ValueTask<VexExportDataSet> FetchAsync(VexQuery query, CancellationToken cancellationToken)
        {
            var claim = new VexClaim(
                "CVE-2025-0001",
                "vendor",
                new VexProduct("pkg:demo/app", "Demo"),
                VexClaimStatus.Affected,
                new VexClaimDocument(VexDocumentFormat.Csaf, "sha256:demo", new Uri("https://example.org/demo")),
                DateTimeOffset.UtcNow,
                DateTimeOffset.UtcNow);

            var consensus = new VexConsensus(
                "CVE-2025-0001",
                claim.Product,
                VexConsensusStatus.Affected,
                DateTimeOffset.UtcNow,
                new[] { new VexConsensusSource("vendor", VexClaimStatus.Affected, "sha256:demo", 1.0) },
                conflicts: null,
                policyVersion: "baseline/v1",
                summary: "affected");

            var dataSet = new VexExportDataSet(
                ImmutableArray.Create(consensus),
                ImmutableArray.Create(claim),
                ImmutableArray.Create("vendor"));

            LastDataSet = dataSet;
            return ValueTask.FromResult(dataSet);
        }
    }

    private sealed class DummyExporter : IVexExporter
    {
        public DummyExporter(VexExportFormat format)
        {
            Format = format;
        }

        public VexExportFormat Format { get; }

        public VexContentAddress Digest(VexExportRequest request)
            => new("sha256", "deadbeef");

        public ValueTask<VexExportResult> SerializeAsync(VexExportRequest request, Stream output, CancellationToken cancellationToken)
        {
            var bytes = System.Text.Encoding.UTF8.GetBytes("{}");
            output.Write(bytes);
            return ValueTask.FromResult(new VexExportResult(new VexContentAddress("sha256", "deadbeef"), bytes.Length, ImmutableDictionary<string, string>.Empty));
        }
    }

}