using System.Buffers.Binary;
using System.Text;

namespace Titano.Video;

public enum VideoCodec
{
    H264,
    Hevc,
}

/// <summary>
/// Multiplexer ISO Base Media File Format (MP4) scritto interamente in-house.
///
/// Struttura generata: ftyp → mdat (in streaming, con dimensione a 64 bit corretta a
/// posteriori) → moov. I campioni vengono scritti sul file di uscita mentre arrivano
/// dall'encoder: nessun file temporaneo, nessuna copia intermedia dell'intero flusso.
///
/// Il muxer accetta bitstream in formato Annex-B (quello prodotto dagli encoder di sistema),
/// ne estrae i parameter set per la configurazione del codec e converte le NAL nel formato
/// a lunghezza prefissata richiesto dal contenitore.
/// </summary>
public sealed class Mp4Muxer : IDisposable
{
    private readonly Stream _output;
    private readonly VideoCodec _codec;
    private readonly int _width;
    private readonly int _height;
    private readonly uint _timescale;

    private readonly List<uint> _sampleSizes = [];
    private readonly List<uint> _sampleDurations = [];
    private readonly List<int> _syncSamples = [];

    private readonly List<byte[]> _vps = [];
    private readonly List<byte[]> _sps = [];
    private readonly List<byte[]> _pps = [];

    private long _mdatHeaderPosition;
    private long _mdatPayloadStart;
    private long _mediaDuration;
    private bool _finished;

    /// <summary>
    /// Buffer di lavoro riusato per la conversione Annex-B → lunghezza prefissata.
    /// Parte piccolo di proposito: così il percorso di crescita viene esercitato da
    /// qualunque sequenza, verifica sintetica compresa, invece di restare latente
    /// fino al primo fotogramma chiave di una ripresa ad alta risoluzione.
    /// </summary>
    private byte[] _scratch = new byte[64 * 1024];

    public int SampleCount => _sampleSizes.Count;
    public long BytesWritten { get; private set; }
    public bool HasParameterSets => _sps.Count > 0;

    public Mp4Muxer(Stream output, VideoCodec codec, int width, int height, uint timescale)
    {
        if (!output.CanSeek) throw new ArgumentException("Il flusso di uscita deve essere posizionabile.", nameof(output));
        _output = output;
        _codec = codec;
        _width = width;
        _height = height;
        _timescale = timescale == 0 ? 90000u : timescale;

        WriteFileTypeBox();
        BeginMediaData();
    }

    /// <summary>
    /// Registra parameter set forniti fuori banda dall'encoder (blob MF_MT_MPEG_SEQUENCE_HEADER).
    /// </summary>
    public void AddParameterSets(ReadOnlySpan<byte> annexB)
    {
        foreach (var nal in AnnexBParser.EnumerateNals(annexB))
        {
            ClassifyAndStore(annexB, nal);
        }
    }

    /// <summary>
    /// Scrive un campione codificato. <paramref name="duration"/> è espressa nella timescale
    /// della traccia, così ogni fotogramma può avere durata propria (playback adattivo).
    /// </summary>
    public void WriteSample(ReadOnlySpan<byte> annexB, uint duration)
    {
        ObjectDisposedException.ThrowIf(_finished, this);

        int written = 0;
        bool keyframe = false;

        foreach (var nal in AnnexBParser.EnumerateNals(annexB))
        {
            var payload = annexB.Slice(nal.Start, nal.Length);
            if (payload.Length == 0) continue;

            if (IsParameterSet(payload[0]))
            {
                ClassifyAndStore(annexB, nal);
                continue;                                   // i parameter set vivono in stsd, non in mdat
            }

            if (IsKeyframeNal(payload[0])) keyframe = true;

            EnsureScratch(written + payload.Length + 4);
            BinaryPrimitives.WriteUInt32BigEndian(_scratch.AsSpan(written), (uint)payload.Length);
            payload.CopyTo(_scratch.AsSpan(written + 4));
            written += payload.Length + 4;
        }

        if (written == 0) return;

        _output.Write(_scratch, 0, written);
        BytesWritten += written;

        if (keyframe) _syncSamples.Add(_sampleSizes.Count + 1);   // gli indici in stss partono da 1
        _sampleSizes.Add((uint)written);
        _sampleDurations.Add(duration);
        _mediaDuration += duration;
    }

    /// <summary>Chiude mdat, scrive moov e finalizza il file.</summary>
    public void Finish()
    {
        if (_finished) return;
        _finished = true;

        long mdatEnd = _output.Position;
        long mdatSize = mdatEnd - _mdatHeaderPosition;

        // Correzione della dimensione a 64 bit riservata all'apertura del box.
        _output.Position = _mdatHeaderPosition + 8;
        Span<byte> size = stackalloc byte[8];
        BinaryPrimitives.WriteInt64BigEndian(size, mdatSize);
        _output.Write(size);

        _output.Position = mdatEnd;
        WriteMovieBox();
        _output.Flush();
    }

    public void Dispose() => Finish();

    // ------------------------------------------------------------------ box di apertura

    private void WriteFileTypeBox()
    {
        using var box = new BoxWriter(_output, "ftyp");
        box.WriteFourCc("isom");
        box.WriteUInt32(0x200);
        box.WriteFourCc("isom");
        box.WriteFourCc("iso2");
        box.WriteFourCc(_codec == VideoCodec.H264 ? "avc1" : "hvc1");
        box.WriteFourCc("mp41");
    }

    private void BeginMediaData()
    {
        _mdatHeaderPosition = _output.Position;
        Span<byte> header = stackalloc byte[16];
        BinaryPrimitives.WriteUInt32BigEndian(header, 1);              // size = 1 → largesize a 64 bit
        Encoding.ASCII.GetBytes("mdat", header[4..]);
        BinaryPrimitives.WriteInt64BigEndian(header[8..], 16);         // valore provvisorio
        _output.Write(header);
        _mdatPayloadStart = _output.Position;
    }

    // ------------------------------------------------------------------ moov

    private void WriteMovieBox()
    {
        uint movieTimescale = 1000;
        long movieDuration = _timescale == 0 ? 0 : _mediaDuration * movieTimescale / _timescale;

        using var moov = new BoxWriter(_output, "moov");

        using (var mvhd = moov.Child("mvhd"))
        {
            mvhd.WriteFullBoxHeader(0, 0);
            mvhd.WriteUInt32(0);                       // creation_time
            mvhd.WriteUInt32(0);                       // modification_time
            mvhd.WriteUInt32(movieTimescale);
            mvhd.WriteUInt32((uint)movieDuration);
            mvhd.WriteUInt32(0x00010000);              // rate 1.0
            mvhd.WriteUInt16(0x0100);                  // volume 1.0
            mvhd.WriteUInt16(0);                       // reserved
            mvhd.WriteUInt32(0);
            mvhd.WriteUInt32(0);
            mvhd.WriteMatrix();
            for (int i = 0; i < 6; i++) mvhd.WriteUInt32(0);  // pre_defined
            mvhd.WriteUInt32(2);                       // next_track_ID
        }

        using (var trak = moov.Child("trak"))
        {
            using (var tkhd = trak.Child("tkhd"))
            {
                tkhd.WriteFullBoxHeader(0, 0x000007);  // enabled | in movie | in preview
                tkhd.WriteUInt32(0);
                tkhd.WriteUInt32(0);
                tkhd.WriteUInt32(1);                   // track_ID
                tkhd.WriteUInt32(0);                   // reserved
                tkhd.WriteUInt32((uint)movieDuration);
                tkhd.WriteUInt32(0);
                tkhd.WriteUInt32(0);
                tkhd.WriteUInt16(0);                   // layer
                tkhd.WriteUInt16(0);                   // alternate_group
                tkhd.WriteUInt16(0);                   // volume (0 per il video)
                tkhd.WriteUInt16(0);
                tkhd.WriteMatrix();
                tkhd.WriteUInt32((uint)_width << 16);
                tkhd.WriteUInt32((uint)_height << 16);
            }

            using var mdia = trak.Child("mdia");

            using (var mdhd = mdia.Child("mdhd"))
            {
                mdhd.WriteFullBoxHeader(0, 0);
                mdhd.WriteUInt32(0);
                mdhd.WriteUInt32(0);
                mdhd.WriteUInt32(_timescale);
                mdhd.WriteUInt32((uint)_mediaDuration);
                mdhd.WriteUInt16(0x55C4);              // lingua "und" impacchettata a 5 bit
                mdhd.WriteUInt16(0);
            }

            using (var hdlr = mdia.Child("hdlr"))
            {
                hdlr.WriteFullBoxHeader(0, 0);
                hdlr.WriteUInt32(0);                   // pre_defined
                hdlr.WriteFourCc("vide");
                hdlr.WriteUInt32(0);
                hdlr.WriteUInt32(0);
                hdlr.WriteUInt32(0);
                hdlr.WriteCString("Titano Video Handler");
            }

            using var minf = mdia.Child("minf");

            using (var vmhd = minf.Child("vmhd"))
            {
                vmhd.WriteFullBoxHeader(0, 1);
                vmhd.WriteUInt16(0);                   // graphicsmode
                vmhd.WriteUInt16(0);                   // opcolor
                vmhd.WriteUInt16(0);
                vmhd.WriteUInt16(0);
            }

            using (var dinf = minf.Child("dinf"))
            using (var dref = dinf.Child("dref"))
            {
                dref.WriteFullBoxHeader(0, 0);
                dref.WriteUInt32(1);                   // entry_count
                using var url = dref.Child("url ");
                url.WriteFullBoxHeader(0, 1);          // flag 1 = dati nello stesso file
            }

            using var stbl = minf.Child("stbl");
            WriteSampleDescription(stbl);
            WriteTimeToSample(stbl);
            WriteSyncSamples(stbl);
            WriteSampleToChunk(stbl);
            WriteSampleSizes(stbl);
            WriteChunkOffsets(stbl);
        }
    }

    private void WriteSampleDescription(BoxWriter stbl)
    {
        using var stsd = stbl.Child("stsd");
        stsd.WriteFullBoxHeader(0, 0);
        stsd.WriteUInt32(1);                           // entry_count

        string entryName = _codec == VideoCodec.H264 ? "avc1" : "hvc1";
        using var entry = stsd.Child(entryName);

        for (int i = 0; i < 6; i++) entry.WriteByte(0); // reserved
        entry.WriteUInt16(1);                           // data_reference_index
        entry.WriteUInt16(0);                           // pre_defined
        entry.WriteUInt16(0);                           // reserved
        for (int i = 0; i < 3; i++) entry.WriteUInt32(0);
        entry.WriteUInt16((ushort)_width);
        entry.WriteUInt16((ushort)_height);
        entry.WriteUInt32(0x00480000);                  // 72 dpi orizzontali
        entry.WriteUInt32(0x00480000);                  // 72 dpi verticali
        entry.WriteUInt32(0);                           // reserved
        entry.WriteUInt16(1);                           // frame_count
        entry.WritePascalString32("Titano");            // compressorname
        entry.WriteUInt16(0x0018);                      // profondità 24 bit
        entry.WriteUInt16(0xFFFF);                      // pre_defined = -1

        if (_codec == VideoCodec.H264) WriteAvcConfiguration(entry);
        else WriteHevcConfiguration(entry);

        using (var colr = entry.Child("colr"))
        {
            colr.WriteFourCc("nclx");
            colr.WriteUInt16(1);                        // primarie BT.709
            colr.WriteUInt16(1);                        // funzione di trasferimento BT.709
            colr.WriteUInt16(1);                        // matrice BT.709
            colr.WriteByte(0);                          // range televisivo (16-235)
        }

        using var pasp = entry.Child("pasp");
        pasp.WriteUInt32(1);                            // pixel quadrati
        pasp.WriteUInt32(1);
    }

    private void WriteAvcConfiguration(BoxWriter entry)
    {
        using var avcc = entry.Child("avcC");
        byte[] sps = _sps.Count > 0 ? _sps[0] : [];

        avcc.WriteByte(1);                                              // configurationVersion
        avcc.WriteByte(sps.Length > 1 ? sps[1] : (byte)0x64);           // AVCProfileIndication
        avcc.WriteByte(sps.Length > 2 ? sps[2] : (byte)0x00);           // profile_compatibility
        avcc.WriteByte(sps.Length > 3 ? sps[3] : (byte)0x28);           // AVCLevelIndication
        avcc.WriteByte(0xFF);                                           // 6 bit riservati + lengthSizeMinusOne = 3
        avcc.WriteByte((byte)(0xE0 | Math.Min(_sps.Count, 31)));        // 3 bit riservati + numOfSPS

        foreach (var set in _sps)
        {
            avcc.WriteUInt16((ushort)set.Length);
            avcc.WriteBytes(set);
        }

        avcc.WriteByte((byte)Math.Min(_pps.Count, 255));
        foreach (var set in _pps)
        {
            avcc.WriteUInt16((ushort)set.Length);
            avcc.WriteBytes(set);
        }
    }

    private void WriteHevcConfiguration(BoxWriter entry)
    {
        using var hvcc = entry.Child("hvcC");
        byte[] sps = _sps.Count > 0 ? _sps[0] : [];

        // profile_tier_level occupa 12 byte subito dopo i due byte di header NAL e il byte
        // che contiene sps_video_parameter_set_id / sps_max_sub_layers_minus1.
        Span<byte> ptl = stackalloc byte[12];
        if (sps.Length >= 15) sps.AsSpan(3, 12).CopyTo(ptl);
        else ptl[0] = 0x01;                                             // Main profile come ripiego

        hvcc.WriteByte(1);                                              // configurationVersion
        hvcc.WriteByte(ptl[0]);                                         // profile_space/tier/profile_idc
        hvcc.WriteBytes(ptl[1..5]);                                     // general_profile_compatibility_flags
        hvcc.WriteBytes(ptl[5..11]);                                    // general_constraint_indicator_flags
        hvcc.WriteByte(ptl[11]);                                        // general_level_idc
        hvcc.WriteUInt16(0xF000);                                       // min_spatial_segmentation = 0
        hvcc.WriteByte(0xFC);                                           // parallelismType sconosciuto
        hvcc.WriteByte(0xFD);                                           // chromaFormat 4:2:0
        hvcc.WriteByte(0xF8);                                           // bitDepthLumaMinus8 = 0
        hvcc.WriteByte(0xF8);                                           // bitDepthChromaMinus8 = 0
        hvcc.WriteUInt16(0);                                            // avgFrameRate (non dichiarato)
        hvcc.WriteByte(0x0F);                                           // costantFrameRate/temporalId + lengthSizeMinusOne
        hvcc.WriteByte((byte)((_vps.Count > 0 ? 1 : 0) + (_sps.Count > 0 ? 1 : 0) + (_pps.Count > 0 ? 1 : 0)));

        WriteHevcArray(hvcc, 32, _vps);
        WriteHevcArray(hvcc, 33, _sps);
        WriteHevcArray(hvcc, 34, _pps);
    }

    private static void WriteHevcArray(BoxWriter hvcc, byte nalType, List<byte[]> sets)
    {
        if (sets.Count == 0) return;
        hvcc.WriteByte((byte)(0x80 | nalType));                         // array_completeness + NAL_unit_type
        hvcc.WriteUInt16((ushort)sets.Count);
        foreach (var set in sets)
        {
            hvcc.WriteUInt16((ushort)set.Length);
            hvcc.WriteBytes(set);
        }
    }

    private void WriteTimeToSample(BoxWriter stbl)
    {
        // Codifica a corse: le durate uguali consecutive occupano una sola voce.
        var runs = new List<(uint Count, uint Delta)>();
        foreach (uint delta in _sampleDurations)
        {
            if (runs.Count > 0 && runs[^1].Delta == delta) runs[^1] = (runs[^1].Count + 1, delta);
            else runs.Add((1, delta));
        }

        using var stts = stbl.Child("stts");
        stts.WriteFullBoxHeader(0, 0);
        stts.WriteUInt32((uint)runs.Count);
        foreach (var (count, delta) in runs)
        {
            stts.WriteUInt32(count);
            stts.WriteUInt32(delta);
        }
    }

    private void WriteSyncSamples(BoxWriter stbl)
    {
        // Se ogni campione è un punto di sincronizzazione la tabella si omette per convenzione.
        if (_syncSamples.Count == _sampleSizes.Count || _syncSamples.Count == 0) return;

        using var stss = stbl.Child("stss");
        stss.WriteFullBoxHeader(0, 0);
        stss.WriteUInt32((uint)_syncSamples.Count);
        foreach (int index in _syncSamples) stss.WriteUInt32((uint)index);
    }

    private void WriteSampleToChunk(BoxWriter stbl)
    {
        using var stsc = stbl.Child("stsc");
        stsc.WriteFullBoxHeader(0, 0);
        stsc.WriteUInt32(1);                            // una sola voce: tutti i campioni in un chunk
        stsc.WriteUInt32(1);                            // first_chunk
        stsc.WriteUInt32((uint)Math.Max(1, _sampleSizes.Count));
        stsc.WriteUInt32(1);                            // sample_description_index
    }

    private void WriteSampleSizes(BoxWriter stbl)
    {
        using var stsz = stbl.Child("stsz");
        stsz.WriteFullBoxHeader(0, 0);
        stsz.WriteUInt32(0);                            // dimensione variabile
        stsz.WriteUInt32((uint)_sampleSizes.Count);
        foreach (uint size in _sampleSizes) stsz.WriteUInt32(size);
    }

    private void WriteChunkOffsets(BoxWriter stbl)
    {
        using var co64 = stbl.Child("co64");
        co64.WriteFullBoxHeader(0, 0);
        co64.WriteUInt32(1);
        co64.WriteUInt64((ulong)_mdatPayloadStart);
    }

    // ------------------------------------------------------------------ NAL

    private bool IsParameterSet(byte header)
    {
        if (_codec == VideoCodec.H264)
        {
            int type = header & 0x1F;
            return type is 7 or 8;                      // SPS, PPS
        }
        int hevcType = (header >> 1) & 0x3F;
        return hevcType is 32 or 33 or 34;              // VPS, SPS, PPS
    }

    private bool IsKeyframeNal(byte header)
    {
        if (_codec == VideoCodec.H264) return (header & 0x1F) == 5;     // IDR
        int type = (header >> 1) & 0x3F;
        return type is >= 16 and <= 21;                 // BLA/IDR/CRA
    }

    private void ClassifyAndStore(ReadOnlySpan<byte> source, AnnexBParser.NalRange nal)
    {
        if (nal.Length <= 0) return;
        var payload = source.Slice(nal.Start, nal.Length);
        byte header = payload[0];

        List<byte[]> target;
        if (_codec == VideoCodec.H264)
        {
            int type = header & 0x1F;
            if (type == 7) target = _sps;
            else if (type == 8) target = _pps;
            else return;
        }
        else
        {
            int type = (header >> 1) & 0x3F;
            if (type == 32) target = _vps;
            else if (type == 33) target = _sps;
            else if (type == 34) target = _pps;
            else return;
        }

        var copy = payload.ToArray();
        foreach (var existing in target)
        {
            if (existing.AsSpan().SequenceEqual(copy)) return;
        }
        target.Add(copy);
    }

    /// <summary>
    /// Allarga il buffer di conversione conservandone il contenuto.
    ///
    /// La copia non è un dettaglio: il buffer accumula le NAL di un campione una dopo l'altra,
    /// quindi rimpiazzarlo con un array nuovo azzererebbe quelle già scritte. Il campione
    /// uscirebbe con un prefisso di lunghezza nullo in testa e nessun decodificatore lo
    /// supererebbe. Il caso si presenta al primo fotogramma chiave, l'unico abbastanza
    /// grande da superare la capacità iniziale.
    /// </summary>
    private void EnsureScratch(int required)
    {
        if (_scratch.Length >= required) return;

        long size = Math.Max(_scratch.Length, 1);
        while (size < required) size *= 2;
        Array.Resize(ref _scratch, (int)Math.Min(size, Array.MaxLength));
    }
}