import XCTest
@testable import Receiver

final class H264DepacketizerTests: XCTestCase {
    private let start: [UInt8] = [0x00, 0x00, 0x01]

    private func packet(seq: Int, ts: Int, payload: [UInt8], marker: Bool = false) -> RtpPacket {
        RtpPacket(header: RtpHeader(sequenceNumber: seq, timestamp: ts, marker: marker), payload: payload)
    }

    func testSingleNalTwoPackets() {
        let d = H264Depacketizer()
        // SPS (type 7) then a slice (type 5), closed by the marker.
        let first = d.depacketize(packet(seq: 1, ts: 100, payload: [0x67, 0xAA, 0xBB]))
        XCTAssertNil(first.accessUnit)
        let second = d.depacketize(packet(seq: 2, ts: 100, payload: [0x41, 0x01, 0x02], marker: true))
        XCTAssertEqual(second.accessUnit, start + [0x67, 0xAA, 0xBB] + start + [0x41, 0x01, 0x02])
        XCTAssertTrue(second.isKeyFrame)
        XCTAssertFalse(second.frameDropped)
    }

    func testFuAReassembly() {
        let d = H264Depacketizer()
        // NAL: header 0x41 (type 1, NRI 2) + payload 0x11 0x22 0x33 0x44.
        // FU indicator = (0x41 & 0xE0) | 28 = 0x5C.
        let indicator: UInt8 = 0x5C
        _ = d.depacketize(packet(seq: 1, ts: 100, payload: [indicator, 0x81, 0x11]))
        _ = d.depacketize(packet(seq: 2, ts: 100, payload: [indicator, 0x01, 0x22]))
        let done = d.depacketize(packet(seq: 3, ts: 100, payload: [indicator, 0x41, 0x33, 0x44], marker: true))
        XCTAssertEqual(done.accessUnit, start + [0x41, 0x11, 0x22, 0x33, 0x44])
        XCTAssertFalse(done.frameDropped)
    }

    func testDropsGappedFrames() {
        let d = H264Depacketizer()
        _ = d.depacketize(packet(seq: 1, ts: 100, payload: [0x41, 0x01]))
        // seq 2 is missing; the frame must be reported dropped, not delivered.
        let tail = d.depacketize(packet(seq: 3, ts: 100, payload: [0x41, 0x02], marker: true))
        XCTAssertNil(tail.accessUnit)
        XCTAssertTrue(tail.frameDropped)
    }

    func testSeparateFramesByMarker() {
        let d = H264Depacketizer()
        let au1 = d.depacketize(packet(seq: 1, ts: 100, payload: [0x41, 0xAA], marker: true))
        XCTAssertEqual(au1.accessUnit, start + [0x41, 0xAA])
        let au2 = d.depacketize(packet(seq: 2, ts: 200, payload: [0x41, 0xBB], marker: true))
        XCTAssertEqual(au2.accessUnit, start + [0x41, 0xBB])
    }

    func testDropsFuWithoutStart() {
        let d = H264Depacketizer()
        // Continuation (no S bit) without any start packet.
        let result = d.depacketize(packet(seq: 1, ts: 100, payload: [0x7C, 0x41, 0x11], marker: true))
        XCTAssertNil(result.accessUnit)
        XCTAssertTrue(result.frameDropped)
    }

    func testDropsStillFragmentedAtMarker() {
        let d = H264Depacketizer()
        _ = d.depacketize(packet(seq: 1, ts: 100, payload: [0x7C, 0x81, 0x11]))
        // Marker arrives while the FU-A NAL is still open.
        let result = d.depacketize(packet(seq: 2, ts: 100, payload: [0x41, 0x01], marker: true))
        XCTAssertNil(result.accessUnit)
        XCTAssertTrue(result.frameDropped)
    }

    func testDropsUnsupportedPacketization() {
        let d = H264Depacketizer()
        let stapA: UInt8 = 24 // STAP-A
        let result = d.depacketize(packet(seq: 1, ts: 100, payload: [stapA, 0x00, 0x05, 0x41, 0x01], marker: true))
        XCTAssertNil(result.accessUnit)
        XCTAssertTrue(result.frameDropped)
    }

    func testDropsTimestampChangeWithoutMarker() {
        let d = H264Depacketizer()
        _ = d.depacketize(packet(seq: 1, ts: 100, payload: [0x41, 0x01]))
        // The stale frame is reported dropped, but this packet starts (and
        // closes) the next frame — matching the C++ depacketizer.
        let result = d.depacketize(packet(seq: 2, ts: 200, payload: [0x41, 0x02], marker: true))
        XCTAssertEqual(result.accessUnit, start + [0x41, 0x02])
        XCTAssertTrue(result.frameDropped)
    }

    func testKeyframeDetectionRequiresParameterSets() {
        let d = H264Depacketizer()
        let plain = d.depacketize(packet(seq: 1, ts: 100, payload: [0x41, 0x01], marker: true))
        XCTAssertFalse(plain.isKeyFrame)

        let d2 = H264Depacketizer()
        let withSps = d2.depacketize(packet(seq: 1, ts: 100, payload: [0x67, 0xAA, 0x88, 0x68, 0xBB, 0x41, 0x01], marker: true))
        XCTAssertTrue(withSps.isKeyFrame)
    }

    func testDropsShortFuPackets() {
        let d = H264Depacketizer()
        // FU-A packet without its FU header byte.
        let result = d.depacketize(packet(seq: 1, ts: 100, payload: [0x7C], marker: true))
        XCTAssertNil(result.accessUnit)
        XCTAssertTrue(result.frameDropped)
    }
}