import asyncio
import logging
import math
from queue import Queue
from struct import pack
from subprocess import Popen
from threading import Thread
from typing import Iterator, List, Tuple
try:
    from aiortc.mediastreams import EncodedMediaStreamTrack
except:
    log = logging.getLogger(__name__)
    log.error("wrong aiortc version")

PACKET_MAX = 1300

NAL_TYPE_FU_A = 28
NAL_TYPE_STAP_A = 24

NAL_HEADER_SIZE = 1
FU_A_HEADER_SIZE = 2
LENGTH_FIELD_SIZE = 2
STAP_A_HEADER_SIZE = NAL_HEADER_SIZE + LENGTH_FIELD_SIZE
logger = logging.getLogger(__name__)

class PiH264CameraOutput(object):
    def __init__(self):
        logger.debug('Spawning background buffer process')
        self.nal_queue = Queue(3)
        self.nal_buffer = None

    def write(self, buf):
        if self.nal_buffer is None:
            self.nal_buffer = buf
        else:
            self.nal_buffer += buf
        if (
                len(self.nal_buffer) < 64
        ):  # Just making sure to pack SPS/PPS within a single buffer
            return
        if not self.nal_queue.full():
            self.nal_queue.put(self.nal_buffer)
        self.nal_buffer = None

    def flush(self):
        logger.debug('Waiting for background conversion process to exit')




class PiH264Relay(EncodedMediaStreamTrack):

    kind = "video"
    _start: float
    _timestamp: int

    def __init__(self, video_rate, clock_rate=90000) -> None:
        super().__init__()
        self.nal_queue = Queue(3)
        self._timestamp = 0
        self._frame_time = 1 / video_rate
        self._clock_rate = clock_rate
        self.nal_buffer = None

    def write(self, buf: bytes):
        if self.nal_buffer is None:
            self.nal_buffer = buf
        else:
            self.nal_buffer += buf
        if (
            len(self.nal_buffer) < 64
        ):  # Just making sure to pack SPS/PPS within a single buffer
            return
        if not self.nal_queue.full():
            self.nal_queue.put(self.nal_buffer)
        self.nal_buffer = None

    @staticmethod
    def _packetize_fu_a(data: bytes) -> List[bytes]:
        available_size = PACKET_MAX - FU_A_HEADER_SIZE
        payload_size = len(data) - NAL_HEADER_SIZE
        num_packets = math.ceil(payload_size / available_size)
        num_larger_packets = payload_size % num_packets
        package_size = payload_size // num_packets

        f_nri = data[0] & (0x80 | 0x60)  # fni of original header
        nal = data[0] & 0x1F

        fu_indicator = f_nri | NAL_TYPE_FU_A

        fu_header_end = bytes([fu_indicator, nal | 0x40])
        fu_header_middle = bytes([fu_indicator, nal])
        fu_header_start = bytes([fu_indicator, nal | 0x80])
        fu_header = fu_header_start

        packages = []
        offset = NAL_HEADER_SIZE
        while offset < len(data):
            if num_larger_packets > 0:
                num_larger_packets -= 1
                payload = data[offset : offset + package_size + 1]
                offset += package_size + 1
            else:
                payload = data[offset : offset + package_size]
                offset += package_size

            if offset == len(data):
                fu_header = fu_header_end

            packages.append(fu_header + payload)

            fu_header = fu_header_middle
        assert offset == len(data), "incorrect fragment data"

        return packages

    @staticmethod
    def _packetize_stap_a(
        data: bytes, packages_iterator: Iterator[bytes]
    ) -> Tuple[bytes, bytes]:
        counter = 0
        available_size = PACKET_MAX - STAP_A_HEADER_SIZE

        stap_header = NAL_TYPE_STAP_A | (data[0] & 0xE0)

        payload = bytes()
        try:
            nalu = data  # with header
            while len(nalu) <= available_size and counter < 9:
                stap_header |= nalu[0] & 0x80

                nri = nalu[0] & 0x60
                if stap_header & 0x60 < nri:
                    stap_header = stap_header & 0x9F | nri

                available_size -= LENGTH_FIELD_SIZE + len(nalu)
                counter += 1
                payload += pack("!H", len(nalu)) + nalu
                nalu = next(packages_iterator)

            if counter == 0:
                nalu = next(packages_iterator)
        except StopIteration:
            nalu = None

        if counter <= 1:
            return data, nalu
        else:
            return bytes([stap_header]) + payload, nalu

    @classmethod
    def _packetize(cls, packages: Iterator[bytes]) -> List[bytes]:
        packetized_packages = []

        packages_iterator = iter(packages)
        package = next(packages_iterator, None)
        while package is not None:
            if len(package) > PACKET_MAX:
                packetized_packages.extend(cls._packetize_fu_a(package))
                package = next(packages_iterator, None)
            else:
                packetized, package = cls._packetize_stap_a(package, packages_iterator)
                packetized_packages.append(packetized)
        return packetized_packages

    @staticmethod
    def _split_bitstream(buf: bytes) -> Iterator[bytes]:
        # TODO: write in a more pytonic way,
        # translate from: https://github.com/aizvorski/h264bitstream/blob/master/h264_nal.c#L134
        i = 0
        while True:
            while (buf[i] != 0 or buf[i + 1] != 0 or buf[i + 2] != 0x01) and (
                buf[i] != 0 or buf[i + 1] != 0 or buf[i + 2] != 0 or buf[i + 3] != 0x01
            ):
                i += 1  # skip leading zero
                if i + 4 >= len(buf):
                    return
            if buf[i] != 0 or buf[i + 1] != 0 or buf[i + 2] != 0x01:
                i += 1
            i += 3
            nal_start = i
            if (buf[nal_start] & 0x1F) != 7 and (
                buf[nal_start] & 0x1F
            ) != 8:  # Assuming only SPS/PPS are packaged with other NALs
                yield buf[nal_start:]
                return
            while (buf[i] != 0 or buf[i + 1] != 0 or buf[i + 2] != 0) and (
                buf[i] != 0 or buf[i + 1] != 0 or buf[i + 2] != 0x01
            ):
                i += 1
                # FIXME: the next line fails when reading a nal that ends
                # exactly at the end of the data
                if i + 3 >= len(buf):
                    nal_end = len(buf)
                    yield buf[nal_start:]
                    return  # did not find nal end, stream ended first
            nal_end = i
            yield buf[nal_start:nal_end]

    async def recv_encoded(self, keyframe=False) -> List[bytes]:
        while True:
            if self.nal_queue.empty():
                await asyncio.sleep(self._frame_time)
                continue
            nal = self.nal_queue.get()
            #if (nal[4] & 0x1F) != 0x01 or not keyframe:
            #    break
            break
        self._timestamp += int(self._frame_time * self._clock_rate)
        timestamp = self._timestamp
        return self._packetize(self._split_bitstream(nal)), timestamp

