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

namespace Jellyfin.Plugin.Enigma2;

/// <summary>Rebase cached/live MPEG-TS segments to one presentation clock.</summary>
public static class TransportStreamClock
{
    private const long ClockMask = (1L << 33) - 1;
    private static int Payload(byte[] data, int start)
    {
        var control = (data[start + 3] >> 4) & 3;
        if ((control & 1) == 0) return -1;
        var payload = start + 4 + ((control & 2) != 0 ? data[start + 4] + 1 : 0);
        return payload < start + 188 ? payload : -1;
    }
    private static long ReadPts(byte[] data, int p) => ((long)(data[p] & 14) << 29) | ((long)data[p + 1] << 22)
        | ((long)(data[p + 2] & 254) << 14) | ((long)data[p + 3] << 7) | ((long)data[p + 4] >> 1);
    private static void WritePts(byte[] data, int p, long pts)
    {
        pts &= ClockMask;
        data[p] = (byte)((data[p] & 240L) | ((pts >> 29) & 14) | 1);
        data[p + 1] = (byte)(pts >> 22); data[p + 2] = (byte)(((pts >> 14) & 254) | 1);
        data[p + 3] = (byte)(pts >> 7); data[p + 4] = (byte)((pts << 1) | 1);
    }
    private static bool Pes(byte[] data, int start, int p) => p >= 0 && (data[start + 1] & 64) != 0 && p + 14 <= start + 188
        && data[p] == 0 && data[p + 1] == 0 && data[p + 2] == 1 && (data[p + 7] & 128) != 0;
    public static byte[] Rebase(byte[] input, long videoStart90k, bool alignAudio = true)
    {
        if (input.Length % 188 != 0) throw new InvalidDataException("Unaligned MPEG-TS segment");
        long? firstVideo = null, firstAudio = null;
        for (var start = 0; start < input.Length; start += 188)
        {
            if (input[start] != 71) throw new InvalidDataException("Invalid MPEG-TS sync");
            var p = Payload(input, start);
            if (!Pes(input, start, p)) continue;
            if (firstVideo == null && input[p + 3] >= 224 && input[p + 3] <= 239) firstVideo = ReadPts(input, p + 9);
            if (firstAudio == null && input[p + 3] >= 192 && input[p + 3] <= 223) firstAudio = ReadPts(input, p + 9);
            if (firstVideo != null && firstAudio != null) break;
        }
        if (firstVideo == null) throw new InvalidDataException("MPEG-TS segment lacks a video presentation timestamp");
        var offset = videoStart90k - firstVideo.Value;
        // Independent AAC encoders and HLS boundaries have different priming
        // offsets. Anchor the first audio PES of each segment as well, so a
        // switch cannot send an audio timestamp backwards by a few milliseconds.
        var audioOffset = alignAudio ? videoStart90k - (firstAudio ?? firstVideo.Value) : offset;
        using var output = new MemoryStream(input.Length + 188 * 6); var pids = new HashSet<int>();
        for (var start = 0; start < input.Length; start += 188)
        {
            var packet = input.AsSpan(start, 188).ToArray(); var pid = ((packet[1] & 31) << 8) | packet[2];
            if (pid != 8191 && pids.Add(pid))
            {
                // Each independently muxed/cached clip may reset continuity
                // counters. Signal that reset in-band, including to FFmpeg's TS
                // demuxer, which does not reset them solely from HLS tags.
                var reset = new byte[188]; Array.Fill(reset, (byte)255);
                reset[0] = 71; reset[1] = (byte)(packet[1] & 31); reset[2] = packet[2];
                var counter = packet[3] & 15;
                if ((packet[3] & 16) != 0) counter = (counter + 15) & 15;
                reset[3] = (byte)(32 | counter); reset[4] = 183; reset[5] = 128;
                output.Write(reset);
            }
            var p = Payload(packet, 0);
            if (Pes(packet, 0, p))
            {
                var shift = packet[p + 3] >= 192 && packet[p + 3] <= 223 ? audioOffset : offset;
                WritePts(packet, p + 9, ReadPts(packet, p + 9) + shift);
                if ((packet[p + 7] & 192) == 192 && p + 19 <= 188) WritePts(packet, p + 14, ReadPts(packet, p + 14) + shift);
            }
            if ((packet[3] & 32) != 0 && packet[4] >= 7 && (packet[5] & 16) != 0)
            {
                var clock = (((long)packet[6] << 25) | ((long)packet[7] << 17) | ((long)packet[8] << 9) | ((long)packet[9] << 1) | ((long)packet[10] >> 7));
                clock = (clock + offset) & ClockMask;
                packet[6] = (byte)(clock >> 25); packet[7] = (byte)(clock >> 17); packet[8] = (byte)(clock >> 9); packet[9] = (byte)(clock >> 1);
                packet[10] = (byte)((packet[10] & 127L) | ((clock & 1) << 7));
            }
            output.Write(packet);
        }
        return output.ToArray();
    }
}