using System;
using System.Collections.Generic;
using System.IO;

/// <summary>
/// Remuxes one HLS MPEG-TS segment (H.264 video + AAC ADTS audio) into a standalone MP4, without
/// re-encoding. Windows Media Foundation (Unity's VideoPlayer on Windows) can't read MPEG-TS, but plays
/// H.264/AAC in MP4. Checked against ffmpeg: decoded video and audio match the source segment exactly.
/// </summary>
public static class TsToMp4
{
    static readonly int[] AacRates = { 96000, 88200, 64000, 48000, 44100, 32000, 24000, 22050, 16000, 12000, 11025, 8000, 7350 };

    struct Pes { public long Pts; public byte[] Data; }
    struct VSample { public long Pts; public byte[] Data; public bool Key; }

    /// <summary>Returns the MP4 bytes, or null if the segment has no usable video.</summary>
    public static byte[] Convert(byte[] ts)
    {
        Demux(ts, out var videoPes, out var audioPes);
        var video = VideoSamples(videoPes, out var sps, out var pps);
        if (video.Count == 0 || sps == null || pps == null) return null;
        var audio = AudioFrames(audioPes, out int aacObject, out int sfIndex, out int channels);
        ParseSpsSize(sps, out int width, out int height);

        // Durations: video from PTS deltas (90 kHz), audio is 1024 samples per AAC frame
        var vDur = new List<uint>();
        for (int i = 0; i < video.Count - 1; i++) vDur.Add((uint)Math.Max(1, video[i + 1].Pts - video[i].Pts));
        uint avg = 3750;
        if (vDur.Count > 0) { long sum = 0; foreach (var d in vDur) sum += d; avg = (uint)(sum / vDur.Count); }
        vDur.Add(avg);
        long vTotal = 0; foreach (var d in vDur) vTotal += d;
        bool hasAudio = audio.Count > 0 && sfIndex >= 0 && sfIndex < AacRates.Length;
        int rate = hasAudio ? AacRates[sfIndex] : 48000;
        long aTotal = hasAudio ? audio.Count * 1024L : 0;

        var ftyp = Box("ftyp", Cat(Ascii("isom"), U32(512), Ascii("isomiso2avc1mp41")));
        long vBytes = 0; foreach (var s in video) vBytes += s.Data.Length;

        byte[] moov = BuildMoov(video, vDur, vTotal, sps, pps, width, height, audio, hasAudio, aacObject, sfIndex, channels, rate, aTotal, 0, 0);
        uint baseOffset = (uint)(ftyp.Length + moov.Length + 8);
        moov = BuildMoov(video, vDur, vTotal, sps, pps, width, height, audio, hasAudio, aacObject, sfIndex, channels, rate, aTotal, baseOffset, baseOffset + (uint)vBytes);

        using (var ms = new MemoryStream())
        {
            ms.Write(ftyp, 0, ftyp.Length);
            ms.Write(moov, 0, moov.Length);
            long mdatLen = 8 + vBytes;
            if (hasAudio) foreach (var a in audio) mdatLen += a.Length;
            var hdr = Cat(U32((uint)mdatLen), Ascii("mdat"));
            ms.Write(hdr, 0, hdr.Length);
            foreach (var s in video) ms.Write(s.Data, 0, s.Data.Length);
            if (hasAudio) foreach (var a in audio) ms.Write(a, 0, a.Length);
            return ms.ToArray();
        }
    }

    // ---------------- MPEG-TS demux ----------------

    static void Demux(byte[] data, out List<Pes> video, out List<Pes> audio)
    {
        video = new List<Pes>(); audio = new List<Pes>();
        int pmtPid = -1, vPid = -1, aPid = -1;
        var open = new Dictionary<int, (long pts, MemoryStream buf)>();

        void Flush(int pid, List<Pes> v, List<Pes> a)
        {
            if (!open.TryGetValue(pid, out var p)) return;
            if (p.buf.Length > 0)
            {
                var pes = new Pes { Pts = p.pts, Data = p.buf.ToArray() };
                if (pid == vPid) v.Add(pes); else a.Add(pes);
            }
            open.Remove(pid);
        }

        for (int off = 0; off + 188 <= data.Length; off += 188)
        {
            if (data[off] != 0x47) continue;
            bool pusi = (data[off + 1] & 0x40) != 0;
            int pid = ((data[off + 1] & 0x1F) << 8) | data[off + 2];
            int afc = (data[off + 3] >> 4) & 3;
            int i = off + 4;
            if (afc == 2 || afc == 3) i += 1 + data[off + 4];
            if (afc == 0 || afc == 2 || i >= off + 188) continue;
            int end = off + 188;

            if (pid == 0 && pusi)
            {
                int t = i + 1 + data[i];
                int secLen = ((data[t + 1] & 0x0F) << 8) | data[t + 2];
                for (int j = t + 8; j < t + 3 + secLen - 4; j += 4)
                {
                    int prog = (data[j] << 8) | data[j + 1];
                    if (prog != 0) pmtPid = ((data[j + 2] & 0x1F) << 8) | data[j + 3];
                }
                continue;
            }
            if (pid == pmtPid && pusi)
            {
                int t = i + 1 + data[i];
                int secLen = ((data[t + 1] & 0x0F) << 8) | data[t + 2];
                int progInfo = ((data[t + 10] & 0x0F) << 8) | data[t + 11];
                int j = t + 12 + progInfo;
                while (j < t + 3 + secLen - 4)
                {
                    int streamType = data[j];
                    int esPid = ((data[j + 1] & 0x1F) << 8) | data[j + 2];
                    int esInfo = ((data[j + 3] & 0x0F) << 8) | data[j + 4];
                    if (streamType == 0x1B && vPid < 0) vPid = esPid;
                    if (streamType == 0x0F && aPid < 0) aPid = esPid;
                    j += 5 + esInfo;
                }
                continue;
            }
            if (pid != vPid && pid != aPid) continue;

            if (pusi)
            {
                Flush(pid, video, audio);
                if (end - i < 9 || data[i] != 0 || data[i + 1] != 0 || data[i + 2] != 1) continue;
                int flags = data[i + 7], hlen = data[i + 8];
                long pts = -1;
                if ((flags & 0x80) != 0)
                {
                    int b = i + 9;
                    pts = ((long)(data[b] & 0x0E) << 29) | ((long)data[b + 1] << 22) | ((long)(data[b + 2] & 0xFE) << 14) | ((long)data[b + 3] << 7) | ((long)data[b + 4] >> 1);
                }
                var ms = new MemoryStream();
                int start = i + 9 + hlen;
                if (start < end) ms.Write(data, start, end - start);
                open[pid] = (pts, ms);
            }
            else if (open.TryGetValue(pid, out var cur))
            {
                cur.buf.Write(data, i, end - i);
            }
        }
        Flush(vPid, video, audio);
        Flush(aPid, video, audio);
    }

    // ---------------- H.264 ----------------

    static List<VSample> VideoSamples(List<Pes> pes, out byte[] sps, out byte[] pps)
    {
        sps = null; pps = null;
        var samples = new List<VSample>();
        foreach (var p in pes)
        {
            var body = new MemoryStream();
            bool key = false;
            foreach (var nal in NalUnits(p.Data))
            {
                int type = nal[0] & 0x1F;
                if (type == 9) continue;              // access unit delimiter
                if (type == 7) sps = nal;
                if (type == 8) pps = nal;
                if (type == 5) key = true;
                var len = U32((uint)nal.Length);
                body.Write(len, 0, 4);
                body.Write(nal, 0, nal.Length);
            }
            if (body.Length > 0) samples.Add(new VSample { Pts = p.Pts, Data = body.ToArray(), Key = key });
        }
        return samples;
    }

    static List<byte[]> NalUnits(byte[] es)
    {
        var units = new List<byte[]>();
        int n = es.Length, i = 0, start = -1;
        while (i + 3 <= n)
        {
            bool sc3 = es[i] == 0 && es[i + 1] == 0 && es[i + 2] == 1;
            bool sc4 = i + 3 < n && es[i] == 0 && es[i + 1] == 0 && es[i + 2] == 0 && es[i + 3] == 1;
            if (sc3 || sc4)
            {
                int len = sc3 ? 3 : 4;
                if (start >= 0 && i > start) units.Add(Slice(es, start, i - start));
                start = i + len; i += len;
            }
            else i++;
        }
        if (start >= 0 && start < n) units.Add(Slice(es, start, n - start));
        return units;
    }

    /// <summary>Width and height from an SPS (with emulation prevention removed and frame cropping applied).</summary>
    static void ParseSpsSize(byte[] spsNal, out int width, out int height)
    {
        width = 1280; height = 720;
        try
        {
            var rbsp = new List<byte>();
            for (int i = 1; i < spsNal.Length; i++)
            {
                if (i + 2 < spsNal.Length && spsNal[i] == 0 && spsNal[i + 1] == 0 && spsNal[i + 2] == 3) { rbsp.Add(0); rbsp.Add(0); i += 2; continue; }
                rbsp.Add(spsNal[i]);
            }
            var r = new BitReader(rbsp.ToArray());
            int profile = r.Bits(8); r.Bits(16);   // constraint flags + level
            r.Ue();                                // seq_parameter_set_id
            int chroma = 1;
            if (profile == 100 || profile == 110 || profile == 122 || profile == 244 || profile == 44 || profile == 83 || profile == 86 || profile == 118 || profile == 128 || profile == 138 || profile == 139 || profile == 134 || profile == 135)
            {
                chroma = r.Ue();
                if (chroma == 3) r.Bits(1);
                r.Ue(); r.Ue(); r.Bits(1);
                if (r.Bits(1) == 1)
                    for (int k = 0; k < (chroma != 3 ? 8 : 12); k++)
                        if (r.Bits(1) == 1)
                        {
                            int size = k < 6 ? 16 : 64, last = 8, next = 8;
                            for (int j = 0; j < size; j++) { if (next != 0) next = (last + r.Se() + 256) % 256; last = next == 0 ? last : next; }
                        }
            }
            r.Ue();                                // log2_max_frame_num_minus4
            int pocType = r.Ue();
            if (pocType == 0) r.Ue();
            else if (pocType == 1) { r.Bits(1); r.Se(); r.Se(); int cycle = r.Ue(); for (int k = 0; k < cycle; k++) r.Se(); }
            r.Ue(); r.Bits(1);                     // max_num_ref_frames, gaps
            int wMbs = r.Ue() + 1, hMapUnits = r.Ue() + 1;
            int frameMbsOnly = r.Bits(1);
            if (frameMbsOnly == 0) r.Bits(1);
            r.Bits(1);                             // direct_8x8_inference
            int cl = 0, cr = 0, ct = 0, cb = 0;
            if (r.Bits(1) == 1) { cl = r.Ue(); cr = r.Ue(); ct = r.Ue(); cb = r.Ue(); }
            int cropX = chroma == 0 ? 1 : 2, cropY = (chroma == 1 ? 2 : 1) * (2 - frameMbsOnly);
            width = wMbs * 16 - (cl + cr) * cropX;
            height = (2 - frameMbsOnly) * hMapUnits * 16 - (ct + cb) * cropY;
        }
        catch { width = 1280; height = 720; }
    }

    class BitReader
    {
        readonly byte[] d; int pos;
        public BitReader(byte[] data) { d = data; }
        public int Bits(int n) { int v = 0; for (int k = 0; k < n; k++) { v = (v << 1) | ((d[pos >> 3] >> (7 - (pos & 7))) & 1); pos++; } return v; }
        public int Ue() { int zeros = 0; while (Bits(1) == 0) zeros++; return (1 << zeros) - 1 + (zeros > 0 ? Bits(zeros) : 0); }
        public int Se() { int k = Ue(); return (k & 1) == 1 ? (k + 1) / 2 : -(k / 2); }
    }

    // ---------------- AAC ----------------

    static List<byte[]> AudioFrames(List<Pes> pes, out int objectType, out int sfIndex, out int channels)
    {
        objectType = 2; sfIndex = -1; channels = 2;
        var frames = new List<byte[]>();
        foreach (var p in pes)
        {
            var es = p.Data; int i = 0;
            while (i + 7 <= es.Length)
            {
                if (es[i] != 0xFF || (es[i + 1] & 0xF0) != 0xF0) { i++; continue; }
                bool protectionAbsent = (es[i + 1] & 1) == 1;
                int profile = (es[i + 2] >> 6) & 3;
                int sfi = (es[i + 2] >> 2) & 0xF;
                int chan = ((es[i + 2] & 1) << 2) | (es[i + 3] >> 6);
                int flen = ((es[i + 3] & 3) << 11) | (es[i + 4] << 3) | (es[i + 5] >> 5);
                int hdr = protectionAbsent ? 7 : 9;
                if (flen < hdr || i + flen > es.Length) break;
                if (sfIndex < 0) { objectType = profile + 1; sfIndex = sfi; channels = chan; }
                frames.Add(Slice(es, i + hdr, flen - hdr));
                i += flen;
            }
        }
        return frames;
    }

    // ---------------- MP4 boxes ----------------

    static byte[] BuildMoov(List<VSample> video, List<uint> vDur, long vTotal, byte[] sps, byte[] pps, int w, int h,
                            List<byte[]> audio, bool hasAudio, int aacObject, int sfIndex, int channels, int rate, long aTotal,
                            uint vOffset, uint aOffset)
    {
        uint vMs = (uint)(vTotal * 1000 / 90000), aMs = hasAudio ? (uint)(aTotal * 1000 / rate) : 0;

        var sync = new List<uint>();
        for (int i = 0; i < video.Count; i++) if (video[i].Key) sync.Add((uint)(i + 1));
        var vSizes = new List<uint>(); foreach (var s in video) vSizes.Add((uint)s.Data.Length);
        var vStbl = Box("stbl", Cat(StsdAvc1(sps, pps, w, h), Stts(vDur),
                                    sync.Count > 0 && sync.Count < video.Count ? FullBox("stss", 0, 0, Cat(U32((uint)sync.Count), U32s(sync))) : new byte[0],
                                    StscOne((uint)video.Count), Stsz(vSizes), Stco(vOffset)));
        var vTrak = Box("trak", Cat(Tkhd(1, vMs, w, h, false),
            Box("mdia", Cat(Mdhd(90000, (uint)vTotal), Hdlr("vide", "VideoHandler"),
                Box("minf", Cat(FullBox("vmhd", 0, 1, new byte[8]), Dinf(), vStbl))))));

        byte[] aTrak = new byte[0];
        if (hasAudio)
        {
            var aDur = new List<uint>(); var aSizes = new List<uint>();
            foreach (var f in audio) { aDur.Add(1024); aSizes.Add((uint)f.Length); }
            var aStbl = Box("stbl", Cat(StsdMp4a(aacObject, sfIndex, channels, rate), Stts(aDur), StscOne((uint)audio.Count), Stsz(aSizes), Stco(aOffset)));
            aTrak = Box("trak", Cat(Tkhd(2, aMs, 0, 0, true),
                Box("mdia", Cat(Mdhd((uint)rate, (uint)aTotal), Hdlr("soun", "SoundHandler"),
                    Box("minf", Cat(FullBox("smhd", 0, 0, new byte[4]), Dinf(), aStbl))))));
        }
        return Box("moov", Cat(Mvhd(Math.Max(vMs, aMs)), vTrak, aTrak));
    }

    static readonly uint[] Matrix = { 0x00010000, 0, 0, 0, 0x00010000, 0, 0, 0, 0x40000000 };

    static byte[] Mvhd(uint durMs) =>
        FullBox("mvhd", 0, 0, Cat(U32(0), U32(0), U32(1000), U32(durMs), U32(0x00010000), U16(0x0100), new byte[10], U32s(Matrix), new byte[24], U32(3)));

    static byte[] Tkhd(uint id, uint durMs, int w, int h, bool audio) =>
        FullBox("tkhd", 0, 3, Cat(U32(0), U32(0), U32(id), U32(0), U32(durMs), new byte[8], U16(0), U16(0), U16((ushort)(audio ? 0x0100 : 0)), U16(0),
                                  U32s(Matrix), U32((uint)w << 16), U32((uint)h << 16)));

    static byte[] Mdhd(uint timescale, uint dur) => FullBox("mdhd", 0, 0, Cat(U32(0), U32(0), U32(timescale), U32(dur), U16(0x55C4), U16(0)));

    static byte[] Hdlr(string type, string name) => FullBox("hdlr", 0, 0, Cat(U32(0), Ascii(type), new byte[12], Ascii(name), new byte[1]));

    static byte[] Dinf() => Box("dinf", FullBox("dref", 0, 0, Cat(U32(1), FullBox("url ", 0, 1, new byte[0]))));

    static byte[] Stts(List<uint> deltas)
    {
        var runs = new List<(uint count, uint delta)>();
        foreach (var d in deltas)
        {
            if (runs.Count > 0 && runs[runs.Count - 1].delta == d) runs[runs.Count - 1] = (runs[runs.Count - 1].count + 1, d);
            else runs.Add((1, d));
        }
        var body = new MemoryStream();
        var n = U32((uint)runs.Count); body.Write(n, 0, 4);
        foreach (var r in runs) { var c = U32(r.count); var dd = U32(r.delta); body.Write(c, 0, 4); body.Write(dd, 0, 4); }
        return FullBox("stts", 0, 0, body.ToArray());
    }

    static byte[] StscOne(uint count) => FullBox("stsc", 0, 0, Cat(U32(1), U32(1), U32(count), U32(1)));
    static byte[] Stsz(List<uint> sizes) => FullBox("stsz", 0, 0, Cat(U32(0), U32((uint)sizes.Count), U32s(sizes)));
    static byte[] Stco(uint offset) => FullBox("stco", 0, 0, Cat(U32(1), U32(offset)));

    static byte[] StsdAvc1(byte[] sps, byte[] pps, int w, int h)
    {
        var avcC = Cat(new byte[] { 1, sps[1], sps[2], sps[3], 0xFF, 0xE1 }, U16((ushort)sps.Length), sps, new byte[] { 1 }, U16((ushort)pps.Length), pps);
        var entry = Cat(new byte[6], U16(1), new byte[16], U16((ushort)w), U16((ushort)h), U32(0x00480000), U32(0x00480000), U32(0), U16(1),
                        new byte[32], U16(0x0018), U16(0xFFFF), Box("avcC", avcC));
        return FullBox("stsd", 0, 0, Cat(U32(1), Box("avc1", entry)));
    }

    static byte[] StsdMp4a(int objectType, int sfIndex, int channels, int rate)
    {
        var asc = U16((ushort)((objectType << 11) | (sfIndex << 7) | (channels << 3)));
        byte[] Desc(byte tag, byte[] body) => Cat(new byte[] { tag, 0x80, 0x80, 0x80, (byte)body.Length }, body);
        var dec = Desc(0x04, Cat(new byte[] { 0x40, 0x15, 0, 0, 0 }, U32(0), U32(0), Desc(0x05, asc)));
        var es = Desc(0x03, Cat(U16(2), new byte[] { 0 }, dec, Desc(0x06, new byte[] { 0x02 })));
        var entry = Cat(new byte[6], U16(1), new byte[8], U16((ushort)channels), U16(16), U16(0), U16(0), U32((uint)rate << 16), FullBox("esds", 0, 0, es));
        return FullBox("stsd", 0, 0, Cat(U32(1), Box("mp4a", entry)));
    }

    // ---------------- byte helpers ----------------

    static byte[] Box(string type, byte[] payload) => Cat(U32((uint)(8 + payload.Length)), Ascii(type), payload);
    static byte[] FullBox(string type, int version, int flags, byte[] payload) => Box(type, Cat(U32((uint)((version << 24) | flags)), payload));
    static byte[] U32(uint v) => new[] { (byte)(v >> 24), (byte)(v >> 16), (byte)(v >> 8), (byte)v };
    static byte[] U16(ushort v) => new[] { (byte)(v >> 8), (byte)v };
    static byte[] U16(int v) => U16((ushort)v);
    static byte[] U32s(IList<uint> values) { var b = new byte[values.Count * 4]; for (int i = 0; i < values.Count; i++) Buffer.BlockCopy(U32(values[i]), 0, b, i * 4, 4); return b; }
    static byte[] Ascii(string s) => System.Text.Encoding.ASCII.GetBytes(s);
    static byte[] Slice(byte[] a, int start, int len) { var b = new byte[len]; Buffer.BlockCopy(a, start, b, 0, len); return b; }
    static byte[] Cat(params byte[][] parts)
    {
        int n = 0; foreach (var p in parts) n += p.Length;
        var r = new byte[n]; int o = 0;
        foreach (var p in parts) { Buffer.BlockCopy(p, 0, r, o, p.Length); o += p.Length; }
        return r;
    }
}