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

import logging
from os import (
    SEEK_SET,
    fstat,
    ftruncate,
    lseek,
    read,
    write,
)
from typing import IO

from typeguard import typechecked
import hexdump

logger = logging.getLogger(__name__)

#MKV is formatted as an EBML file (Extended Binary Markup Langage).
# cf http://matroska-org.github.io/libebml/specs.html
# It is a Type, Length, Value (TLV) kind of binary file.
# Types are encoded as follows:
# 1xxx xxxx                                  - Class A IDs (2^7 -1 possible values)
# 01xx xxxx  xxxx xxxx                       - Class B IDs (2^14-1 possible values)
# 001x xxxx  xxxx xxxx  xxxx xxxx            - Class C IDs (2^21-1 possible values)
# 0001 xxxx  xxxx xxxx  xxxx xxxx  xxxx xxxx - Class D IDs (2^28-1 possible values)
# Lengths are encoded as follows:
# 1xxx xxxx
# value 0 to  2^7-2
# 01xx xxxx  xxxx xxxx
# value 0 to 2^14-2
# 001x xxxx  xxxx xxxx  xxxx xxxx
# value 0 to 2^21-2
# 0001 xxxx  xxxx xxxx  xxxx xxxx  xxxx xxxx
# value 0 to 2^28-2
# 0000 1xxx  xxxx xxxx  xxxx xxxx  xxxx xxxx  xxxx xxxx
# value 0 to 2^35-2
# 0000 01xx  xxxx xxxx  xxxx xxxx  xxxx xxxx  xxxx xxxx  xxxx xxxx
# value 0 to 2^42-2
# 0000 001x  xxxx xxxx  xxxx xxxx  xxxx xxxx  xxxx xxxx  xxxx xxxx  xxxx xxxx
# value 0 to 2^49-2
# 0000 0001  xxxx xxxx  xxxx xxxx  xxxx xxxx  xxxx xxxx  xxxx xxxx  xxxx xxxx  xxxx xxxx
# value 0 to 2^56-2

@typechecked
def get_ebml_length(length:int) -> bytes|None:
    if 0 <= length <= 2**7-2:
        size = 1
    elif length <= 2**14-2:
        size = 2
    elif length <= 2**21-2:
        size = 3
    elif length <= 2**28-2:
        size = 4
    elif length <= 2**35-2:
        size = 5
    elif length <= 2**42-2:
        size = 6
    elif length <= 2**49-2:
        size = 7
    elif length <= 2**56-2:
        size = 8
    elif length < 0:
        logger.error('Impossible to encode a negative length with EBML.')
        return None
    else:
        logger.error('Impossible to encode a length larger than 2^56-2 with EBML.')
        return None

    encoded_length = length + ((128>>(size-1))<<((size-1)*8))
    res = (encoded_length).to_bytes(size, byteorder='big')
    return res


@typechecked
def change_ebml_element_size(input_file: IO[bytes], position:int, addendum:int) -> int:
    initial_position = position
    infd = input_file.fileno()
    lseek(infd, position, SEEK_SET)

    buf = read(infd, 1)
    element_type = int.from_bytes(buf, byteorder='big')
    mask=128
    found = False
    for i in range(1,5):
        if element_type&mask:
            type_size = i
            found = True
            break
        mask = mask>>1

    if not found:
        logger.error('Size of element type cannot be determined: %d', element_type)
        raise ValueError()

    # We seek to size
    position+=type_size
    lseek(infd, position, SEEK_SET)

    buf = read(infd, 1)
    size_head = int.from_bytes(buf, byteorder='big')
    logger.info('First byte of size: %x', size_head)
    mask=128
    found = False
    for i in range(1,9):
        if size_head&mask:
            size_of_data_size = i
            found = True
            break
        mask = mask>>1

    if not found:
        logger.error('Size of data size cannot be determined: %d', size_head)
        raise ValueError()
    logger.info('Size of data size: %d.', size_of_data_size)

    lseek(infd, position, SEEK_SET)
    old_size_buf = read(infd, size_of_data_size)
    max_size = 2**(size_of_data_size*7)-2
    size_of_data = int.from_bytes(old_size_buf, byteorder='big')
    logger.info('Size of data with mask: %x mask: %d.', size_of_data, mask)
    size_of_data-= (mask<<((size_of_data_size-1)*8))
    logger.info('Found element at position: %d, size of type: %d  size of data: %d \
                maximal size: %d.', initial_position, type_size, size_of_data, max_size)

    new_size = size_of_data+addendum
    delta = 0
    if  new_size > max_size:
        # TODO: Test this code ...
        new_encoded_size = get_ebml_length(new_size)
        if new_encoded_size is None:
            raise ValueError()
        size_of_new_encoded_size = len(new_encoded_size)
        if size_of_new_encoded_size <= size_of_data_size:
            logger.error('New encoded size is smaller (%d) or equal than previous size (%d).\
                         This should not happen.', size_of_new_encoded_size, size_of_data_size)
            raise ValueError()
        # The difference of length between old size field and new one.
        delta = size_of_new_encoded_size - size_of_data_size
        file_length = fstat(infd).st_size
        # We seek after actual length field
        lseek(infd, position+size_of_data_size, SEEK_SET)
        # We read the rest of file
        tail = read(infd, file_length-(position+size_of_data_size))
        # We increase file length
        ftruncate(infd, file_length+delta)
        # We go to the beginning of length field
        lseek(infd, position, SEEK_SET)
        # We write the new length field
        write(infd, new_encoded_size)
        # We overwrite the rest of file with its previous content that has been offset.
        write(infd, tail)
    else:
        size = new_size + ((128>>(size_of_data_size-1))<<((size_of_data_size-1)*8))
        new_size_buf = (size).to_bytes(size_of_data_size, byteorder='big')

        logger.info('Old encoded size: %s New encoded size: %s', hexdump.dump(old_size_buf,sep=':'),
                    hexdump.dump(new_size_buf, sep=':'))
        lseek(infd, position, SEEK_SET)
        write(infd, new_size_buf)

    # We return the potential increase in size of the file if the length field had to be increased.
    return delta