#!/usr/bin/env python3
'''A module to remove parts of video (.e.g advertisements) with single frame precision.'''

# Standard modules
import json
import logging
import os.path
import re
from datetime import timedelta
from enum import IntEnum, unique
from io import BytesIO, TextIOWrapper
from math import ceil, floor, log
from os import (
    SEEK_SET,
    close,
    fstat,
    ftruncate,
    lseek,
    memfd_create,
    read,
    set_inheritable,
    write,
)
from subprocess import PIPE, Popen
from sys import exit
from typing import IO

# Third party libraries
import hexdump
from iso639 import Lang
from iso639.exceptions import InvalidLanguageValue
from tqdm import tqdm
from typeguard import typechecked

from tscut.h264.avc import (
    get_avc_config_from_h264,
    parse_codec_private
)
from tscut.matroska.ebml import change_ebml_element_size
from tscut.matroska.codec import dump_codec_private_data
from tscut.tools.ffprobe import (
    get_format,
    get_frame_rate,
    get_video_dimensions,
    get_movie_duration,
    get_streams,
    with_subtitles,
    find_subtitles_tracks,
    get_nearest_iframe,
    get_ts_frame
)
from tscut.tools.ffmpeg import (
    extract_all_streams,
    ffmpeg_convert
)
from tscut.tools.timeframe import (
    compare_time_interval,
    parse_time_interval
)
from tscut.tools.discovery import check_required_tools
from tscut.ocr.subtitles import (
    do_ocr,
    extract_srt,
    get_tesseract_supported_lang
)


# Useful SPS/PPS discussion.
# https://copyprogramming.com/howto/including-sps-and-pps-in-a-raw-h264-track
# https://gitlab.com/mbunkus/mkvtoolnix/-/issues/2390

# New strategy: a possible way of handling multiple SPS/PPS gracefully.
# Encode each head and trailer with FFMPEG using only I-frame (to be sure the NAL unit will never
# refer to another image).
# Encode using an different SPS-ID all of them (using sps-id parameter of libx264 library, e.g
# 1 instead of 0).
# For the video track produce only a raw H264 file and a file containing timestamps of the
# different frames.
# For the rest of the tracks (audio, subtitles) produce directly a MKV (this is already done).
# Concatenate all raw H264 in a giant one (like cat), and the same for timestamps of video frames
# (to keep sound and video synchronized).
# Then use mkvmerge to remux the H264 track and the rest of tracks.
# MKVmerge "concatenate" subcommand is able to concatenate different SPS/PPS data into a bigger
# Private Codec Data.
# However, this is proved to be not reliable. Sometimes it results in a AVC context containing
# a single SPS/PPS.
# So we have to rely on a manual parsing of the H264 AVC context of original movie
# and the ones produced for headers and trailers, and then merging them into a bigger AVC context.
# Then finally, change the Private Codec Data in the final MKV.







@unique
class SupportedFormat(IntEnum):
    TS = 1
    MP4 = 2
    MATROSKA = 3

    def __str__(self):
        match self:
            case SupportedFormat.TS:
                return 'mpegts'
            case SupportedFormat.MP4:
                return 'mov,mp4,m4a,3gp,3g2,mj2'
            case SupportedFormat.MATROSKA:
                return 'matroska,webm'
            case _:
                return 'Unsupported format'

# Extract SPS/PPS
# https://gitlab.com/mbunkus/mkvtoolnix/-/issues/2390
# ffmpeg -i <InputFile (before concatenation)> -c:v copy -an -sn -bsf:v trace_headers -t 0.01\
#        -report -loglevel 0 -f null -

# 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()
        out = 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 out.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

    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()
        out = 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 out.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)
    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

@typechecked
def concatenate_h264_parts(h264parts: list[IO[bytes]], output: IO[bytes]) -> None:
    logger = logging.getLogger(__name__)

    total_length = 0
    for h264 in h264parts:
        fd = h264.fileno()
        total_length += fstat(fd).st_size

    logger.info('Total length: %d', total_length)

    outfd = output.fileno()
    lseek(outfd, 0, SEEK_SET)

    pb = tqdm(total=total_length, unit='bytes', desc='Concatenation')
    for h264 in h264parts:
        fd = h264.fileno()
        lseek(fd, 0, SEEK_SET)
        while True:
            buf = read(fd, 1000000)
            if buf is None or len(buf) == 0:
                break
            pos = 0
            while pos < len(buf):
                nb_bytes = write(outfd, buf[pos:])
                pb.update(nb_bytes)
                pos += nb_bytes

def concatenate_h264_ts_parts(h264_ts_parts: list[IO[bytes]], output: IO[bytes]) -> None:
    logger = logging.getLogger(__name__)
    header = '# timestamp format v2\n'

    output.write(header)

    last = 0.
    first = True
    for part in h264_ts_parts:
        if first:
            offset = last
        else:
            # TODO: take framerate into account
            offset = last + 40
        logger.debug('Parsing file: %s. Offset=%d', part, offset)
        isheader = part.readline()
        if (not isheader) or (isheader != header):
            logger.error('Impossible to find a valid header: "%s"', isheader)
            exit(-1)
        while True:
            line = part.readline()
            if not line:
                break
            ts = offset + float(line)
            last = max(last,ts)
            output.write(f'{ts:f}\n')
        if first:
            first = False

# TODO: finish this procedure
def do_coarse_processing(ffmpeg_path:str, ffprobe_path:str, mkvmerge_path:str,
                         input_file: IO[bytes], begin, end, nb_frames, framerate,
                         files_prefix, streams, width, height, temporaries, dump_mem_fd) -> None:
    # pylint: disable=W0613
    logger = logging.getLogger(__name__)

    # Internal video with all streams (video, audio and subtitles)
    internal_mkv_name = f'{files_prefix}.mkv'

    try:
        internal_mkv = open(internal_mkv_name, 'wb+')
    except OSError:
        logger.error('Impossible to create file: %s', internal_mkv_name)
        exit(-1)

    # Extract internal part of MKV
    extract_mkv_part(mkvmerge_path=mkvmerge_path, input_file=input_file, output_file=internal_mkv,
                     begin=begin, end=end)

    temporaries.append(internal_mkv)