# SPDX-License-Identifier: GPL-2.0-or-later
#
# Copyright (C) 2026 Frédéric Tronel

""" All the following code is a transposition of documents:
    ISO/IEC H.264-201602
    ISO/IEC 14496-15
"""

import logging
from math import floor, log

from typeguard import typechecked
import hexdump


logger = logging.getLogger(__name__)

@typechecked
def read_bit(buf:bytes, bit_position: int) -> tuple[int, int]:
    """
    Read a single bit from a byte buffer.

    This function is part of the implementation of the H.264/AVC video compression standard,
    as specified in ISO/IEC H.264-201602 and ISO/IEC 14496-15.
    It takes a byte buffer and a bit position as input, and returns a tuple containing
    the updated bit position and the value of the bit at the specified position.

    Args:
        buf (bytes): The byte buffer to read from.
        bit_position (int): The position of the bit to read, starting from 0.

    Returns:
        tuple[int, int]: A tuple containing the updated bit position (bit_position + 1) and the
        value of the bit (0 or 1).
    """
    byte_position = floor(floor(bit_position/8))
    byte = buf[byte_position]
    bit = (byte >> (7-(bit_position % 8))) & 1
    return bit_position+1, bit

@typechecked
def read_boolean(buf:bytes, bit_position: int) -> tuple[int, bool]:
    """
    Read a boolean value from a byte buffer.

    This function reads a single bit from the byte buffer at the specified position and interprets
    it as a boolean value.
    It returns a tuple containing the updated bit position and the boolean value.

    Args:
        buf (bytes): The byte buffer to read from.
        bit_position (int): The position of the bit to read, starting from 0.

    Returns:
        tuple[int, bool]: A tuple containing the updated bit position and the boolean value 
        (True if the bit is 1, False if the bit is 0).
    """
    bit_position, b = read_bit(buf, bit_position)
    return bit_position, b==1

@typechecked
def read_bits(buf:bytes, bit_position: int, nb_bits: int)  -> tuple[int, int]:
    v = 0
    for _ in range(nb_bits):
        bit_position, bit = read_bit(buf, bit_position)
        v = v*2+bit
    return bit_position, v

@typechecked
def read_byte(buf:bytes, bit_position: int)  -> tuple[int, int]:
    bit_position, b = read_bits(buf, bit_position, 8)
    return bit_position, b

@typechecked
def read_word(buf:bytes, bit_position: int) -> tuple[int, int]:
    bit_position, w = read_bits(buf, bit_position, 16)
    return bit_position, w

@typechecked
def read_long(buf:bytes, bit_position: int) -> tuple[int, int]:
    bit_position, value = read_bits(buf, bit_position, 32)
    return bit_position, value

@typechecked
def read_unsigned_exp_golomb(buf:bytes, bit_position: int) -> tuple[int, int]:
    nb_zeroes=0
    while True:
        bit_position, b = read_bit(buf, bit_position)
        if b!=0:
            break
        nb_zeroes+=1

    v1 = 1
    bit_position, v2 = read_bits(buf, bit_position, nb_zeroes)
    v = (v1<<nb_zeroes)+v2
    return bit_position, v-1

@typechecked
def read_signed_exp_golomb(buf:bytes, bit_position: int) -> tuple[int, int]:
    bit_position, v = read_unsigned_exp_golomb(buf, bit_position)
    match v%2:
        case 0:
            return bit_position, -(v>>1)
        case 1:
            return bit_position, (v+1)>>1
        case _:
            raise AssertionError("unreachable")

@typechecked
def write_bit(buf:bytearray, bit_position: int, b) -> int:
    buf_length = len(buf)
    byte_position = floor(bit_position/8)

    if byte_position >= buf_length:
        extension = bytearray(byte_position+1-buf_length)
        buf.extend(extension)

    buf[byte_position] |= (b<<(7-(bit_position % 8)))
    bit_position+=1

    return bit_position

@typechecked
def write_boolean(buf:bytearray, bit_position: int, b: bool) -> int:
    if b:
        bit_position = write_bit(buf, bit_position, 1)
    else:
        bit_position = write_bit(buf, bit_position, 0)
    return bit_position

@typechecked
def write_bits(buf:bytearray, bit_position: int, v, size) -> int:
    for i in range(size-1,-1,-1):
        b = (v>>i)&1
        bit_position = write_bit(buf, bit_position, b)

    return bit_position

@typechecked
def write_byte(buf:bytearray, bit_position: int, v) -> int:
    bit_position = write_bits(buf, bit_position, v, 8)
    return bit_position

@typechecked
def write_word(buf:bytearray, bit_position: int, v) -> int:
    bit_position = write_bits(buf, bit_position, v, 16)
    return bit_position

@typechecked
def write_long(buf:bytearray, bit_position: int, v) -> int:
    bit_position = write_bits(buf, bit_position, v, 32)
    return bit_position

@typechecked
def write_unsigned_exp_golomb(buf:bytearray, bit_position: int, v) -> int:
    n = floor(log(v+1)/log(2))+1
    # Write zeroes
    bit_position = write_bits(buf, bit_position, 0, n-1)
    bit_position = write_bit(buf, bit_position, 1)
    bit_position = write_bits(buf, bit_position, v+1, n-1)

    return bit_position

@typechecked
def write_signed_exp_golomb(buf:bytearray, bit_position: int, v) -> int:
    if v <= 0:
        bit_position = write_unsigned_exp_golomb(buf, bit_position, -v*2)
    else:
        bit_position = write_unsigned_exp_golomb(buf, bit_position, v*2-1)

    return bit_position

@typechecked
def parse_rbsp_trailing_bits(buf:bytes, bit_position: int) -> int:
    bit_position, one = read_bit(buf, bit_position)
    if one==0:
        raise ValueError(f'Stop bit should be equal to one. Read: {one:d}')
    while bit_position%8 != 0:
        bit_position, zero = read_bit(buf, bit_position)
        if zero==1:
            raise ValueError('Trailing bit should be equal to zero')

    return bit_position

@typechecked
def write_rbsp_trailing_bits(buf:bytearray, bit_position: int) -> int:
    bit_position = write_bit(buf, bit_position, 1)
    while bit_position%8 != 0:
        bit_position = write_bit(buf, bit_position, 0)

    return bit_position

@typechecked
def more_rbsp_data(buf:bytes, bit_position: int) -> bool:
    logger.debug('Is there more data in buffer of length: %d at bit position: %d',
                 len(buf), bit_position)

    byte_length = len(buf)
    bit_length = byte_length*8

    # We are at the end of buffer
    if bit_position == bit_length:
        return False

    found = False
    for i in range(bit_length-1,-1,-1):
        pos, b = read_bit(buf, i)
        if b == 1:
            found = True
            break

    if not found:
        raise ValueError('Impossible to find trailing stop bit !')

    # No more data
    if bit_position == pos:
        return False

    return True

# Convert from RBSP (Raw Byte Sequence Payload)  to SODB (String Of Data Bits)
@typechecked
def rbsp_to_sodb(buf:bytes) -> bytes:
    logger.debug('RBSP: %s', hexdump.dump(buf, sep=':'))

    res = buf
    for b in [ b'\x00', b'\x01', b'\x02', b'\x03']:
        pattern = b'\x00\x00\x03'+b
        replacement = b'\x00\x00' + b
        res = res.replace(pattern, replacement)

    logger.debug('SODB: %s', hexdump.dump(res, sep=':'))
    return res

# Reverse operation SODB to RBSP.
@typechecked
def sodb_to_rbsp(buf:bytes) -> bytes:
    logger.debug('SODB: %s', hexdump.dump(buf, sep=':'))

    res = buf
    for b in [ b'\x03', b'\x00', b'\x01', b'\x02']:
        pattern = b'\x00\x00'+b
        replacement = b'\x00\x00\x03' + b
        res = res.replace(pattern, replacement)

    logger.debug('RBSP: %s', hexdump.dump(res, sep=':'))
    return res

# Useful for SPS and PPS
@typechecked
def parse_scaling_list(buf:bytes, bit_position: int, size) -> tuple[int,list[int]]:
    res = []
    last_scale = 8
    next_scale = 8
    for _ in range(size):
        if next_scale != 0:
            bit_position, delta_scale = read_signed_exp_golomb(buf, bit_position)
            next_scale = (last_scale+delta_scale+256) % 256
        v = last_scale if next_scale==0 else next_scale
        res.append(v)
        last_scale = v

    return bit_position,res

# TODO: test optimized version.
# The ISO/IEC H.264-201602 seems to take into account the case where the end of the deltas list
#  is full of zeroes.
@typechecked
def write_scaling_list(buf:bytearray, bit_position: int, size, matrix:list[int],
                       optimized: bool = False) -> int:
    logger.debug('Dumping matrix: %s of size: %d, size parameter: %d.', matrix, len(matrix), size)

    prev = 8
    deltas = []
    for i in range(size):
        v = matrix[i]
        delta = v - prev
        deltas.append(delta)
        prev = v

    if not optimized:
        for delta in deltas:
            bit_position = write_signed_exp_golomb(buf, bit_position, delta)
    else:
        logger.error('Not yet implemented')
        exit(-1)
        # reverse = deltas.reverse()
        # compressed = False
        # while len(reverse)>0:
        #     if reverse[0] == 0:
        #         compressed = True
        #         reverse.pop()
        #     else:
        #         break
        # deltas = reverse.reverse()
        # if compressed:
        #     deltas.append(0)
        # for delta in deltas:
        #     bit_position = write_signed_exp_golomb(buf, bit_position, delta)

    return bit_position