// Claude slop... im not skilled amount to do this bullshit

use bytes::Bytes;

pub struct VideoFrame {
    pub data: Bytes,
    pub is_keyframe: bool,
    pub timestamp_ms: u32,
}

pub struct AudioFrame {
    pub data: Bytes,
    pub timestamp_ms: u32,
}

pub struct H264Parser {
    sps: Option<Vec<u8>>,
    pps: Option<Vec<u8>>,
}

impl H264Parser {
    pub fn new() -> Self {
        Self {
            sps: None,
            pps: None,
        }
    }

    /// Parse an RTMP VideoDataReceived payload. Returns None for sequence
    /// header packets (which carry SPS/PPS but no displayable frame).
    pub fn parse(&mut self, bytes: &[u8], timestamp_ms: u32) -> Option<VideoFrame> {
        if bytes.len() < 5 {
            return None;
        }

        let frame_type = (bytes[0] >> 4) & 0x0F;
        let codec_id = bytes[0] & 0x0F;

        if codec_id != 7 {
            return None; // not H.264
        }

        let avc_packet_type = bytes[1];
        // bytes[2..5] are the composition time offset — not needed for sending
        let payload = &bytes[5..];

        match avc_packet_type {
            0 => {
                self.parse_sequence_header(payload);
                None
            }
            1 => {
                let is_keyframe = frame_type == 1;
                let data = self.avcc_to_annexb(payload, is_keyframe)?;
                Some(VideoFrame {
                    data: Bytes::from(data),
                    is_keyframe,
                    timestamp_ms,
                })
            }
            _ => None,
        }
    }

    fn parse_sequence_header(&mut self, payload: &[u8]) {
        // AVCDecoderConfigurationRecord layout:
        // [0] configurationVersion
        // [1] AVCProfileIndication
        // [2] profile_compatibility
        // [3] AVCLevelIndication
        // [4] 0xFF (lower 2 bits = lengthSizeMinusOne, always 3 meaning 4-byte lengths)
        // [5] 0xE0 | numSPS
        // [6..] SPS entries: 2-byte length + bytes
        // then: numPPS, PPS entries: 2-byte length + bytes
        if payload.len() < 7 {
            return;
        }

        let mut i = 5;

        let num_sps = (payload[i] & 0x1F) as usize;
        i += 1;

        for _ in 0..num_sps {
            if i + 2 > payload.len() {
                return;
            }
            let len = u16::from_be_bytes([payload[i], payload[i + 1]]) as usize;
            i += 2;
            if i + len > payload.len() {
                return;
            }
            self.sps = Some(payload[i..i + len].to_vec());
            i += len;
        }

        if i >= payload.len() {
            return;
        }

        let num_pps = payload[i] as usize;
        i += 1;

        for _ in 0..num_pps {
            if i + 2 > payload.len() {
                return;
            }
            let len = u16::from_be_bytes([payload[i], payload[i + 1]]) as usize;
            i += 2;
            if i + len > payload.len() {
                return;
            }
            self.pps = Some(payload[i..i + len].to_vec());
            i += len;
        }
    }

    fn avcc_to_annexb(&self, payload: &[u8], is_keyframe: bool) -> Option<Vec<u8>> {
        let mut out = Vec::new();

        // Prepend SPS+PPS before every keyframe so str0m's packetizer
        // can bundle them into a STAP-A alongside the IDR NALU.
        if is_keyframe {
            if let (Some(sps), Some(pps)) = (&self.sps, &self.pps) {
                out.extend_from_slice(&[0, 0, 0, 1]);
                out.extend_from_slice(sps);
                out.extend_from_slice(&[0, 0, 0, 1]);
                out.extend_from_slice(pps);
            }
        }

        // Convert each length-prefixed NALU to an Annex B start-code NALU.
        let mut i = 0;
        while i + 4 <= payload.len() {
            let nalu_len = u32::from_be_bytes(payload[i..i + 4].try_into().unwrap()) as usize;
            i += 4;
            if i + nalu_len > payload.len() {
                break;
            }
            out.extend_from_slice(&[0, 0, 0, 1]);
            out.extend_from_slice(&payload[i..i + nalu_len]);
            i += nalu_len;
        }

        if out.is_empty() { None } else { Some(out) }
    }
}