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


import logging
from io import TextIOWrapper
import re
from datetime import timedelta
from os import (
    read,
    SEEK_SET,
    lseek,
    write,
    fstat,
    set_inheritable,
    ftruncate
)
import os
from subprocess import PIPE, Popen
from typing import IO

from typeguard import typechecked
import hexdump
from tqdm import tqdm

from tscut.matroska.ebml import change_ebml_element_size

# Found codec private data using mkvinfo
@typechecked
def get_codec_private_data_from_mkv(mkvinfo_path:str,
                                    input_file: IO[bytes]) -> tuple[int, bytes]|tuple[None,None]:
    logger = logging.getLogger(__name__)

    infd = input_file.fileno()
    lseek(infd, 0, SEEK_SET)
    set_inheritable(infd, True)
    found = False
    env = {**os.environ, 'LANG': 'C'}
    # Output example
    # Codec's private data: size 48 (H.264 profile: High @L4.0) hexdump 01 64 00 28 ff e1 00 1b 67\
    # 64 00 28 ac d9 40 78 04 4f dc d4 04 04 05 00 00 92 ef 00 1d ad a6 1f 16 2d 96 01 00 06 68 fb\
    # a3 cb 22 c0 fd f8 f8 00 at 406 size 51 data size 48

    with Popen([mkvinfo_path, '-z', '-X', '-P', f'/proc/self/fd/{infd:d}'], stdout=PIPE,
               close_fds=False, env=env) as mkvinfo:
        out, _ = mkvinfo.communicate()
        lines = out.decode('utf8')
        reg_exp = (r"^.*Codec's private data: size ([0-9]+) \(H.264.*\) hexdump "
                  r"(?P<hexdump>([0-9a-f]{2} )+)at (?P<position>[0-9]+) size (?P<size>[0-9]+).*$")
        p = re.compile(reg_exp)
        for line in lines.splitlines():
            m = p.match(line)
            if m is not None:
                size = int(m.group('size'))
                position = int(m.group('position'))
                logger.debug("Found codec private data at position: %s, size: %d", position, size)
                found = True
                mkvinfo.wait()
                break

    if found:
        lseek(infd, position, SEEK_SET)
        data = read(infd, size)
        return position, data

    logger.error("Impossible to retrieve codec private data from MKV !")
    return None, None

@typechecked
def parse_mkv_tree(mkvinfo_path:str, input_file: IO[bytes]) -> dict[str,tuple[int,int]]:
    logger = logging.getLogger(__name__)

    infd = input_file.fileno()
    lseek(infd, 0, SEEK_SET)
    set_inheritable(infd, True)
    env = {**os.environ, 'LANG': 'C'}
    elements = {}

    with Popen([mkvinfo_path, '-z', '-X', '-P', f'/proc/self/fd/{infd:d}'], stdout=PIPE,
               close_fds=False, env=env) as mkvinfo:
        out, _ = mkvinfo.communicate()
        lines = out.decode('utf8')
        prefix = []
        reg_exp = (r"(^(?P<root>\+)|(\|(?P<depth>[ ]*\+))).*at (?P<position>[0-9]+)"
                  r" size (?P<size>[0-9]+).*$")
        p = re.compile(reg_exp)
        prev_depth = -1
        for line in lines.splitlines():
            m = p.match(line)
            if m is None:
                logger.error("Impossible to match line: %s", line)
            else:
                position = int(m.group('position'))
                size = int(m.group('size'))
                root = m.group('root') is not None
                if root:
                    depth = 0
                else:
                    depth = len(m.group('depth'))

                if depth > prev_depth:
                    for _ in range(depth-prev_depth):
                        prefix.append(1)
                elif depth == prev_depth:
                    subid = prefix[-1]
                    subid+=1
                    prefix.pop()
                    prefix.append(subid)
                else:
                    for _ in range(prev_depth-depth):
                        prefix.pop()
                    subid = prefix[-1]
                    subid+=1
                    prefix.pop()
                    prefix.append(subid)

                prev_depth = depth
                key=".".join(map(str, prefix))
                elements[key] = (position, size)

    mkvinfo.wait()
    return elements

@typechecked
def change_codec_private_data(mkvinfo_path:str, input_file: IO[bytes], codec_data:bytes) -> None:
    logger = logging.getLogger(__name__)

    infd = input_file.fileno()
    lseek(infd, 0, SEEK_SET)

    current_length = fstat(infd).st_size
    logger.info('Current size of file: %d', current_length)
    position, current_data = get_codec_private_data_from_mkv(mkvinfo_path, input_file)
    if position is None or current_data is None:
        return None
    current_data_length = len(current_data)
    future_length = current_length - current_data_length + len(codec_data)
    logger.info('Expected size of file: %d', future_length)

    logger.info('Current data at position %d: %s', position, hexdump.dump(current_data, sep=":"))
    logger.info('Future data: %s', hexdump.dump(codec_data, sep=":"))

    elements = parse_mkv_tree(mkvinfo_path, input_file)

    found = False
    for key, (pos,size) in elements.items():
        if pos == position:
            logger.info('Codec private data key: %s', key)
            found = True
            break

    if not found:
        logger.error('Impossible to retrieve the key of codec private data')
        exit(-1)

    if current_length < future_length:
        lseek(infd, position+current_data_length, SEEK_SET)
        tail = read(infd, current_length-(position+current_data_length))
        # We extend the file at the end with zeroes
        ftruncate(infd, future_length)
        lseek(infd, position+len(codec_data), SEEK_SET)
        write(infd, tail)
        lseek(infd, position, SEEK_SET)
        write(infd, codec_data)
    elif current_length == future_length:
        # Almost nothing to do except overwriting old private codec data with new ones.
        lseek(infd, position, SEEK_SET)
        write(infd, codec_data)
    else:
        lseek(infd, position+current_data_length, SEEK_SET)
        tail = read(infd, current_length-(position+current_data_length))
        lseek(infd, position+len(codec_data), SEEK_SET)
        write(infd, tail)
        lseek(infd, position, SEEK_SET)
        write(infd, codec_data)
        # We reduce the length of file.
        ftruncate(infd, future_length)

    # We have to modify the tree elements up to the root that contains the codec private data.
    keys = key.split('.')
    logger.info(keys)

    delta = future_length-current_length
    # if there is no modification of the private codec data, no need to change anything.
    if delta != 0:
        for _ in range(len(keys)-1):
            keys.pop()
            key=".".join(map(str, keys))
            pos, size = elements[key]
            logger.info('Trying to fix element with key: %s at position: %d with actual size: %d.',
                        key, pos, size)
            # Changing an element can increase its size (in very rare case).
            # In that case, we update the new delta that will be larger (because the element has
            # been resized).
            delta+=change_ebml_element_size(input_file, pos, delta)

@typechecked
def extract_mkv_part(mkvmerge_path:str, input_file:IO[bytes], output_file:IO[bytes],
                     begin:timedelta, end:timedelta) -> None:
    logger = logging.getLogger(__name__)

    logger.info('Extract video between I-frames at %s and %s', begin,end)
    infd = input_file.fileno()
    outfd = output_file.fileno()
    lseek(infd, 0, SEEK_SET)
    lseek(outfd, 0, SEEK_SET)
    set_inheritable(infd, True)
    set_inheritable(outfd, True)
    env = {**os.environ, 'LANG': 'C'}
    warnings = []

    command = [mkvmerge_path, '-o', f'/proc/self/fd/{outfd:d}', '--split', f'parts:{begin}-{end}',
               f'/proc/self/fd/{infd:d}']
    logger.debug('Executing: %s', command)

    with Popen(command, stdout=PIPE, close_fds=False, env=env) as mkvmerge:
        pb = tqdm(TextIOWrapper(mkvmerge.stdout, encoding="utf-8"), total=100, unit='%',
                  desc='Extraction')
        for line in pb:
            if line.startswith('Progress :'):
                p = re.compile('^Progress : (?P<progress>[0-9]{1,3})%$')
                m = p.match(line)
                if m is None:
                    logger.error('Impossible to parse progress')
                pb.update(int(m['progress'])-pb.n)
            elif line.startswith('Warning'):
                warnings.append(line)
        pb.update(100-pb.n)
        pb.refresh()
        pb.close()

        status = mkvmerge.wait()
        if status == 1:
            logger.warning('Extraction returns warning')
            for w in warnings:
                logger.warning(w)
        elif status == 2:
            logger.error('Extraction returns errors')


# Merge a list of mkv files passed as input, and produce a new MKV as output
@typechecked
def merge_mkvs(mkvmerge_path:str, inputs: list[IO[bytes]], output_name:str,
               concatenate: bool=True,
               timestamps: dict[int, IO[str]] | None = None) -> IO[bytes]|None:
    logger = logging.getLogger(__name__)

    if timestamps is None:
        timestamps = {}

    fds = []
    try:
        out = open(output_name, 'wb+')
    except OSError:
        logger.error('Impossible to create file: %s', output_name)
        return None

    outfd = out.fileno()
    lseek(outfd, 0, SEEK_SET)
    fds.append(outfd)
    set_inheritable(outfd, True)

    # Timestamps of merged tracks are modified by the length of the preceding track.
    # The default mode ('file') is using the largest timestamp of the whole file which may create
    # desynchronize video and sound.
    merge_params = [mkvmerge_path, '--append-mode', 'track']

    first = True
    partnum = 0
    for mkv in inputs:
        if mkv is not None:
            fd = mkv.fileno()
            fds.append(fd)
            set_inheritable(fd, True)
            # If we pass a timestamps file associated with the considered track, use it.
            if partnum in timestamps:
                tsfd = timestamps[partnum].fileno()
                lseek(tsfd, 0, SEEK_SET)
                fds.append(tsfd)
                set_inheritable(tsfd, True)
                merge_params.extend(['--timestamps', f'{partnum:d}:/proc/self/fd/{tsfd:d}'])
            if first:
                merge_params.append(f'/proc/self/fd/{fd:d}')
                first = False
            elif concatenate:
                merge_params.append(f'+/proc/self/fd/{fd:d}')
            else:
                merge_params.append(f'/proc/self/fd/{fd:d}')
            partnum+=1

    merge_params.extend(['-o', f'/proc/self/fd/{outfd:d}'])

    # We merge all files.
    warnings = []
    env = {**os.environ, 'LANG': 'C'}
    logger.debug('Executing: LANG=C %s', merge_params)

    with Popen(merge_params, stdout=PIPE, close_fds=False, env=env) as mkvmerge:
        pb = tqdm(TextIOWrapper(mkvmerge.stdout, encoding="utf-8"), total=100, unit='%',
                  desc='Merging')
        for line in pb:
            if line.startswith('Progress :'):
                p = re.compile('^Progress : (?P<progress>[0-9]{1,3})%$')
                m = p.match(line)
                if m is None:
                    logger.error('Impossible to parse progress')
                pb.n = int(m['progress'])
                pb.update()
            elif line.startswith('Warning'):
                warnings.append(line)

        status = mkvmerge.wait()
        if status == 1:
            logger.warning('Extraction returns warning')
            for w in warnings:
                logger.warning(w)
        elif status == 2:
            logger.error('Extraction returns errors')

    for fd in fds:
        set_inheritable(fd, False)

    return out



@typechecked
def extract_track_from_mkv(mkvextract_path: str, input_file: IO[bytes], index,
                           output_file: IO[bytes], timestamps) -> None:
    logger = logging.getLogger(__name__)

    infd = input_file.fileno()
    lseek(infd, 0, SEEK_SET)
    set_inheritable(infd, True)

    outfd = output_file.fileno()
    lseek(outfd, 0, SEEK_SET)
    set_inheritable(outfd, True)

    tsfd = timestamps.fileno()
    lseek(tsfd, 0, SEEK_SET)
    set_inheritable(tsfd, True)

    params = [ mkvextract_path, f'/proc/self/fd/{infd:d}', 'tracks',
              f'{index:d}:/proc/self/fd/{outfd:d}', 'timestamps_v2',
              f'{index:d}:/proc/self/fd/{tsfd:d}']

    env = {**os.environ, 'LANG': 'C'}
    logger.debug('Executing: LANG=C %s', params)

    with Popen(params, stdout=PIPE, close_fds=False, env=env) as extract:
        pb = tqdm(TextIOWrapper(extract.stdout, encoding="utf-8"), total=100, unit='%',
                  desc='Extraction of track')
        for line in pb:
            if line.startswith('Progress :'):
                p = re.compile('^Progress : (?P<progress>[0-9]{1,3})%$')
                m = p.match(line)
                if m is None:
                    logger.error('Impossible to parse progress')
                pb.update(int(m['progress'])-pb.n)
        pb.update(100-pb.n)
        pb.refresh()
        pb.close()

    extract.wait()

    if extract.returncode != 0:
        logger.error('Mkvextract returns an error code: %d', extract.returncode)
    else:
        logger.info('Track %d was succesfully extracted.', index)



@typechecked
def remove_video_tracks_from_mkv(mkvmerge_path:str, input_file: IO[bytes],
                                 output_file: IO[bytes]) -> None:
    logger = logging.getLogger(__name__)

    outfd = output_file.fileno()
    infd = input_file.fileno()
    lseek(infd, 0, SEEK_SET)
    lseek(outfd, 0, SEEK_SET)
    set_inheritable(infd, True)
    set_inheritable(outfd, True)

    params = [ mkvmerge_path, '-o', f'/proc/self/fd/{outfd:d}', '-D', f'/proc/self/fd/{infd:d}']
    logger.debug('Executing: LANG=C %s', params)

    env = {**os.environ, 'LANG': 'C'}
    with Popen(params, stdout=PIPE, close_fds=False, env=env) as remove:
        pb = tqdm(TextIOWrapper(remove.stdout, encoding="utf-8"), total=100, unit='%',
                  desc='Removal of video track:')
        for line in pb:
            if line.startswith('Progress :'):
                p = re.compile('^Progress : (?P<progress>[0-9]{1,3})%$')
                m = p.match(line)
                if m is None:
                    logger.error('Impossible to parse progress')
                pb.update(int(m['progress'])-pb.n)
        pb.update(100-pb.n)
        pb.refresh()
        pb.close()

    remove.wait()

    if remove.returncode != 0:
        logger.error('Mkvmerge returns an error code: %d', remove.returncode)
    else:
        logger.info('Video tracks were succesfully extracted.')


@typechecked
def remux_srt_subtitles(mkvmerge_path:str, input_file: IO[bytes], output_filename: str,
                        subtitles) -> None:
    logger = logging.getLogger(__name__)

    try:
        out = open(output_filename, 'w', encoding='utf8')
    except OSError:
        logger.error('Impossible to create file: %s', output_filename)
        return None

    outfd = out.fileno()
    infd = input_file.fileno()
    lseek(infd, 0, SEEK_SET)
    set_inheritable(infd, True)
    set_inheritable(outfd, True)

    mkv_merge_params = [mkvmerge_path, f'/proc/self/fd/{infd:d}']
    for fd, lang in subtitles:
        lseek(fd, 0, SEEK_SET)
        set_inheritable(fd, True)
        mkv_merge_params.extend(['--language', f'0:{lang}', f'/proc/self/fd/{fd:d}'])

    mkv_merge_params.extend(['-o', f'/proc/self/fd/{outfd:d}'])

    warnings = []
    env = {**os.environ, 'LANG': 'C'}
    logger.info('Remux subtitles: %s', mkv_merge_params)
    with Popen(mkv_merge_params, stdout=PIPE, close_fds=False, env=env) as mkvmerge:
        pb = tqdm(TextIOWrapper(mkvmerge.stdout, encoding="utf-8"), total=100, unit='%',
                  desc='Remux subtitles:')
        for line in pb:
            if line.startswith('Progress :'):
                p = re.compile('^Progress : (?P<progress>[0-9]{1,3})%$')
                m = p.match(line)
                if m is None:
                    logger.error('Impossible to parse progress')
                pb.n = int(m['progress'])
                pb.update()
            elif line.startswith('Warning'):
                warnings.append(line)

        status = mkvmerge.wait()
        if status == 1:
            logger.warning('Remux subtitles returns warning')
            for w in warnings:
                logger.warning(w)
        elif status == 2:
            logger.error('Remux subtitles returns errors')

    return None