using System.Text.Json;
using Confluent.Kafka;
using Microsoft.Extensions.Options;
using Minio.DataModel.Args;

public sealed class DataLakeWriterService : BackgroundService
{
    private readonly KafkaOptions _kafkaOptions;
    private readonly DataLakeOptions _opts;
    private readonly MinioOptions _minioOpts;
    private readonly ILogger<DataLakeWriterService> _logger;

    // Buffer key: identifies one Parquet file-to-be.
    // Events sharing a topic, date, and Kafka partition land in the same file.
    private record BufferKey(string Topic, string DatePath, int Partition);

    private record BufferedEvent(string Payload, long Offset);

    private readonly Dictionary<BufferKey, List<BufferedEvent>> _buffer = new();
    // Track the highest offset per topic-partition for post-flush commit.
    private readonly Dictionary<TopicPartition, TopicPartitionOffset> _highWatermarks = new();

    public DataLakeWriterService(
        IOptions<KafkaOptions> kafkaOptions,
        IOptions<DataLakeOptions> opts,
        IOptions<MinioOptions> minioOpts,
        ILogger<DataLakeWriterService> logger)
    {
        _kafkaOptions = kafkaOptions.Value;
        _opts         = opts.Value;
        _minioOpts    = minioOpts.Value;
        _logger       = logger;
    }

    protected override async Task ExecuteAsync(CancellationToken ct)
    {
        var consumerConfig = new ConsumerConfig
        {
            BootstrapServers = _kafkaOptions.BootstrapServers,
            GroupId          = "data-lake-writer",
            AutoOffsetReset  = AutoOffsetReset.Earliest,
            EnableAutoCommit = false,
        }.ApplySecurity(_kafkaOptions);

        using var consumer = new ConsumerBuilder<string, string>(consumerConfig).Build();
        consumer.Subscribe(new[]
        {
            _kafkaOptions.Topics.ObservationRecorded,
            _kafkaOptions.Topics.AlertGenerated,
            _kafkaOptions.Topics.EncounterStatusChanged,
        });

        _logger.LogInformation(
            "DataLakeWriterService started. FlushCount={FlushCount} FlushIntervalSeconds={FlushInterval}",
            _opts.FlushCount, _opts.FlushIntervalSeconds);

        var lastFlush = DateTimeOffset.UtcNow;

        try
        {
            while (!ct.IsCancellationRequested)
            {
                ConsumeResult<string, string>? result;
                try
                {
                    result = consumer.Consume(TimeSpan.FromMilliseconds(500));
                }
                catch (ConsumeException ex)
                {
                    _logger.LogError(ex, "DataLakeWriter consume error");
                    continue;
                }

                if (result is not null)
                    AddToBuffer(result);

                var totalBuffered    = _buffer.Values.Sum(v => v.Count);
                var shouldFlushCount = totalBuffered >= _opts.FlushCount;
                var shouldFlushTime  = DateTimeOffset.UtcNow - lastFlush
                                       >= TimeSpan.FromSeconds(_opts.FlushIntervalSeconds);

                if ((shouldFlushCount || shouldFlushTime) && totalBuffered > 0)
                {
                    await FlushAsync(consumer, ct);
                    lastFlush = DateTimeOffset.UtcNow;
                }
            }
        }
        finally
        {
            // Final flush on shutdown so buffered events are not lost.
            // Use CancellationToken.None — host shutdown cancels the execute token before
            // MinIO uploads finish, which surfaces as TaskCanceledException noise.
            if (_buffer.Values.Sum(v => v.Count) > 0)
            {
                try { await FlushAsync(consumer, CancellationToken.None); }
                catch (Exception ex)
                {
                    _logger.LogError(ex, "DataLakeWriter shutdown flush failed — some events may be re-read on next start");
                }
            }
            consumer.Close();
        }
    }

    private void AddToBuffer(ConsumeResult<string, string> result)
    {
        var datePath  = ExtractDatePath(result.Topic, result.Message.Value);
        var key       = new BufferKey(result.Topic, datePath, result.Partition.Value);

        if (!_buffer.TryGetValue(key, out var list))
        {
            list = new List<BufferedEvent>();
            _buffer[key] = list;
        }
        list.Add(new BufferedEvent(result.Message.Value, result.Offset.Value));

        // Track highest offset per topic-partition for post-flush commit.
        var tp = new TopicPartition(result.Topic, result.Partition);
        _highWatermarks[tp] = new TopicPartitionOffset(tp, result.Offset + 1);
    }

    private async Task FlushAsync(IConsumer<string, string> consumer, CancellationToken ct)
    {
        var flushedKeys = new HashSet<BufferKey>();
        var failedPartitions = new HashSet<int>();

        foreach (var (key, events) in _buffer)
        {
            if (events.Count == 0) continue;

            try
            {
                var firstOffset = events.Min(e => e.Offset);
                var objectKey   = BuildObjectKey(key, firstOffset);
                var bytes       = await BuildParquetAsync(key.Topic, events, key.Partition);

                await UploadToMinioAsync(objectKey, bytes, ct);
                flushedKeys.Add(key);

                _logger.LogInformation(
                    "[DATA-LAKE] Wrote {Count} events → {ObjectKey} ({Bytes} bytes)",
                    events.Count, objectKey, bytes.Length);
            }
            catch (Exception ex)
            {
                failedPartitions.Add(key.Partition);

                if (ex is OperationCanceledException && ct.IsCancellationRequested)
                {
                    _logger.LogInformation(
                        "[DATA-LAKE] Flush canceled for key {Key} during shutdown", key);
                }
                else
                {
                    _logger.LogError(ex, "[DATA-LAKE] Failed to write file for key {Key}", key);
                }
            }
        }

        // Only commit offsets for topic-partitions that had no failures.
        var safeOffsets = _highWatermarks
            .Where(kv => !failedPartitions.Contains(kv.Key.Partition))
            .Select(kv => kv.Value)
            .ToList();

        if (safeOffsets.Count > 0)
        {
            consumer.Commit(safeOffsets);
            _logger.LogInformation(
                "[DATA-LAKE] Committed offsets for {PartitionCount} partitions after flushing {FileCount} files",
                safeOffsets.Count, flushedKeys.Count);
        }

        if (failedPartitions.Count > 0)
        {
            _logger.LogWarning(
                "[DATA-LAKE] Retained buffers for {FailedCount} failed partitions — will retry next flush",
                failedPartitions.Count);
        }

        // Clear only successfully flushed keys; retain failed ones for retry.
        foreach (var key in flushedKeys)
            _buffer.Remove(key);

        // Clear watermarks only for partitions with no failures.
        foreach (var tp in _highWatermarks.Keys.ToList())
        {
            if (!failedPartitions.Contains(tp.Partition))
                _highWatermarks.Remove(tp);
        }
    }

    private string ExtractDatePath(string topic, string payload) =>
        DataLakeEventParser.ExtractDatePath(topic, payload, _kafkaOptions.Topics);

    private async Task<byte[]> BuildParquetAsync(
        string topic, List<BufferedEvent> events, int partition)
    {
        if (topic == _kafkaOptions.Topics.ObservationRecorded)
        {
            var rows = events.Select(e => DataLakeEventParser.ParseObservationRow(
                e.Payload, e.Offset, partition)).ToList();
            return await ParquetFileBuilder.BuildObservationsAsync(rows);
        }
        if (topic == _kafkaOptions.Topics.AlertGenerated)
        {
            var rows = events.Select(e => DataLakeEventParser.ParseAlertRow(
                e.Payload, e.Offset, partition)).ToList();
            return await ParquetFileBuilder.BuildAlertsAsync(rows);
        }
        if (topic == _kafkaOptions.Topics.EncounterStatusChanged)
        {
            var rows = events.Select(e => DataLakeEventParser.ParseEncounterRow(
                e.Payload, e.Offset, partition)).ToList();
            return await ParquetFileBuilder.BuildEncountersAsync(rows);
        }
        throw new InvalidOperationException($"Unknown topic: {topic}");
    }

    // ---------------------------------------------------------------------------
    // Object key and date partition helpers
    // ---------------------------------------------------------------------------

    // File path: observations/2025/01/15/partition-0-offset-0000001000.parquet
    // The date comes from the event timestamp, not the wall clock.
    // Events from the same encounter that cross midnight are written into the date
    // bucket matching their recorded_at timestamp — consistent with how Athena and
    // Spark partition-prune by event time, not ingest time.
    private string BuildObjectKey(BufferKey key, long firstOffset)
    {
        var folder = key.Topic switch
        {
            var t when t == _kafkaOptions.Topics.ObservationRecorded    => "observations",
            var t when t == _kafkaOptions.Topics.AlertGenerated         => "alerts",
            var t when t == _kafkaOptions.Topics.EncounterStatusChanged => "encounters",
            _                                                              => "unknown",
        };
        return $"{folder}/{key.DatePath}/partition-{key.Partition}-offset-{firstOffset:D10}.parquet";
    }

    // ---------------------------------------------------------------------------
    // MinIO upload
    // ---------------------------------------------------------------------------

    private async Task UploadToMinioAsync(string objectKey, byte[] bytes, CancellationToken ct)
    {
        var client = MinioClientFactory.Build(_minioOpts);
        var bucket = _opts.BucketName;

        var exists = await client.BucketExistsAsync(
            new BucketExistsArgs().WithBucket(bucket), ct);
        if (!exists)
            await client.MakeBucketAsync(new MakeBucketArgs().WithBucket(bucket), ct);

        await client.PutObjectAsync(new PutObjectArgs()
            .WithBucket(bucket)
            .WithObject(objectKey)
            .WithStreamData(new MemoryStream(bytes))
            .WithObjectSize(bytes.Length)
            .WithContentType("application/octet-stream"),
            ct);
    }
}