using System.Collections.Immutable;
using System.IO;
using System.Linq;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;
using StellaOps.Vexer.Core;
using StellaOps.Vexer.Policy;
using StellaOps.Vexer.Storage.Mongo;

namespace StellaOps.Vexer.Export;

public interface IExportEngine
{
    ValueTask<VexExportManifest> ExportAsync(VexExportRequestContext context, CancellationToken cancellationToken);
}

public sealed record VexExportRequestContext(
    VexQuery Query,
    VexExportFormat Format,
    DateTimeOffset RequestedAt,
    bool ForceRefresh = false);

public interface IVexExportDataSource
{
    ValueTask<VexExportDataSet> FetchAsync(VexQuery query, CancellationToken cancellationToken);
}

public sealed record VexExportDataSet(
    ImmutableArray<VexConsensus> Consensus,
    ImmutableArray<VexClaim> Claims,
    ImmutableArray<string> SourceProviders);

public sealed class VexExportEngine : IExportEngine
{
    private readonly IVexExportStore _exportStore;
    private readonly IVexPolicyEvaluator _policyEvaluator;
    private readonly IVexExportDataSource _dataSource;
    private readonly IReadOnlyDictionary<VexExportFormat, IVexExporter> _exporters;
    private readonly ILogger<VexExportEngine> _logger;
    private readonly IVexCacheIndex? _cacheIndex;
    private readonly IReadOnlyList<IVexArtifactStore> _artifactStores;

    public VexExportEngine(
        IVexExportStore exportStore,
        IVexPolicyEvaluator policyEvaluator,
        IVexExportDataSource dataSource,
        IEnumerable<IVexExporter> exporters,
        ILogger<VexExportEngine> logger,
        IVexCacheIndex? cacheIndex = null,
        IEnumerable<IVexArtifactStore>? artifactStores = null)
    {
        _exportStore = exportStore ?? throw new ArgumentNullException(nameof(exportStore));
        _policyEvaluator = policyEvaluator ?? throw new ArgumentNullException(nameof(policyEvaluator));
        _dataSource = dataSource ?? throw new ArgumentNullException(nameof(dataSource));
        _logger = logger ?? throw new ArgumentNullException(nameof(logger));
        _cacheIndex = cacheIndex;
        _artifactStores = artifactStores?.ToArray() ?? Array.Empty<IVexArtifactStore>();

        if (exporters is null)
        {
            throw new ArgumentNullException(nameof(exporters));
        }

        _exporters = exporters.ToDictionary(x => x.Format);
    }

    public async ValueTask<VexExportManifest> ExportAsync(VexExportRequestContext context, CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(context);
        var signature = VexQuerySignature.FromQuery(context.Query);

        if (!context.ForceRefresh)
        {
            var cached = await _exportStore.FindAsync(signature, context.Format, cancellationToken).ConfigureAwait(false);
            if (cached is not null)
            {
                _logger.LogInformation("Reusing cached export for {Signature} ({Format})", signature.Value, context.Format);
                return new VexExportManifest(
                    cached.ExportId,
                    cached.QuerySignature,
                    cached.Format,
                    cached.CreatedAt,
                    cached.Artifact,
                    cached.ClaimCount,
                    cached.SourceProviders,
                    fromCache: true,
                    cached.ConsensusRevision,
                    cached.Attestation,
                    cached.SizeBytes);
            }
        }
        else if (_cacheIndex is not null)
        {
            await _cacheIndex.RemoveAsync(signature, context.Format, cancellationToken).ConfigureAwait(false);
            _logger.LogInformation("Force refresh requested; invalidated cache entry for {Signature} ({Format})", signature.Value, context.Format);
        }

        var dataset = await _dataSource.FetchAsync(context.Query, cancellationToken).ConfigureAwait(false);
        var exporter = ResolveExporter(context.Format);

        var exportRequest = new VexExportRequest(
            context.Query,
            dataset.Consensus,
            dataset.Claims,
            context.RequestedAt);

        var digest = exporter.Digest(exportRequest);

        await using var buffer = new MemoryStream();
        var result = await exporter.SerializeAsync(exportRequest, buffer, cancellationToken).ConfigureAwait(false);

        if (_artifactStores.Count > 0)
        {
            var writtenBytes = buffer.ToArray();
            try
            {
                var artifact = new VexExportArtifact(
                    result.Digest,
                    context.Format,
                    writtenBytes,
                    result.Metadata);

                foreach (var store in _artifactStores)
                {
                    await store.SaveAsync(artifact, cancellationToken).ConfigureAwait(false);
                }

                _logger.LogInformation("Stored export artifact {Digest} via {StoreCount} store(s)", result.Digest.ToUri(), _artifactStores.Count);
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Failed to store export artifact {Digest}", result.Digest.ToUri());
                throw;
            }
        }

        var exportId = FormattableString.Invariant($"exports/{context.RequestedAt:yyyyMMddTHHmmssfffZ}/{digest.Digest}");
        var manifest = new VexExportManifest(
            exportId,
            signature,
            context.Format,
            context.RequestedAt,
            digest,
            dataset.Claims.Length,
            dataset.SourceProviders,
            fromCache: false,
            consensusRevision: _policyEvaluator.Version,
            attestation: null,
            sizeBytes: result.BytesWritten);

        await _exportStore.SaveAsync(manifest, cancellationToken).ConfigureAwait(false);

        _logger.LogInformation(
            "Export generated for {Signature} ({Format}) size={SizeBytes} bytes",
            signature.Value,
            context.Format,
            result.BytesWritten);

        return manifest;
    }

    private IVexExporter ResolveExporter(VexExportFormat format)
        => _exporters.TryGetValue(format, out var exporter)
            ? exporter
            : throw new InvalidOperationException($"No exporter registered for format '{format}'.");
}

public static class VexExportServiceCollectionExtensions
{
    public static IServiceCollection AddVexExportEngine(this IServiceCollection services)
    {
        services.AddSingleton<IExportEngine, VexExportEngine>();
        services.AddVexExportCacheServices();
        return services;
    }
}
