'''
List of backends used by the 
core module to download animes.
'''

import os
import re
import time
import copy

import shutil
import requests
import threading
from ffmpeg_progress_yield import FfmpegProgress

import logging

from typing import Callable, Literal
from nekosama import consts
from nekosama import utils

logger = logging.getLogger(__name__)

# Get the FFMPEG path
FFMPEG = shutil.which('ffmpeg')

BACKENDS = [
    'ffmpeg',
    'thread',
    'thread_ffmpeg',
    'safe'
]


METHOD_TYPE = Literal['ffmpeg', 'thread', 'thread_ffmpeg', 'safe']

def reach(method: str) -> Callable:
    '''
    Reach a backend function.
    '''
    
    if not method in BACKENDS:
        raise Exception('Invalid backend', method)
    
    if FFMPEG is None and 'ffmpeg' in method:
        logger.warning('FFMPEG is not installed, falling back')
        method = 'safe'
    
    return eval('bk_' + method)


def bk_ffmpeg(raw: str,
              path: str,
              callback: Callable = None,
              quiet: bool = False,
              **kwargs) -> None:
    '''
    Download using the ffmpeg download feature.
    '''

    # Get the total count
    lenght = len(re.findall(consts.re.fragments, raw))

    # Write the m3u8 file for ffmpeg
    with open('temp.m3u', 'w') as file: file.write(raw)
    
    # Create the FFMPEG command
    command = [
        FFMPEG,
        '-protocol_whitelist', 'file,http,https,tcp,tls,crypto',
        '-i', 'temp.m3u',
        '-acodec', 'copy',
        '-vcodec', 'copy',
        path, '-y'
    ]
    
    for progress in FfmpegProgress(command).run_command_with_progress():
        if callback:
            callback(progress)
    
    # Delete the m3u8 file
    os.remove('temp.m3u')

def dl_frag(req: requests.PreparedRequest,
            timeout: float,
            session: requests.Session,
            data: dict = None,
            todo: list = None) -> bytes | None:
    '''
    Download a single fragment.
    '''
    
    raw = session.send(req).content
    
    while b'<html>' in raw:
        
        # If request fails, put it back in todo
        if todo is not None:
            return todo.append(req)
        
        time.sleep(timeout)
        raw = session.send(req).content
    
    else:
        if data is None: return raw
        key = int(re.findall(consts.re.frag_index, req.url)[0])
        data[key] = raw

def bk_base_thread(raw: str,
                   timeout: float = .05,
                   callback: Callable = None,
                   quiet: bool = False,
                   **kwargs) -> dict[int, bytes]:
    '''
    Download each fragment using threads.
    Theorically faster.
    Instead of writing it to a path, returns it.
    '''
    
    chunks = re.findall(consts.re.fragments, raw)
    reqs = [requests.Request('GET', chunk, consts.headers).prepare() for chunk in chunks]
    session = requests.session()
    
    data = {}
    todo = copy.deepcopy(reqs)
    
    while len(todo):
        current = todo.pop(0)
        
        t = threading.Thread(
            target = dl_frag,
            args = [current, timeout, session, data, todo]
        )
        
        t.start()
        
        time.sleep(timeout)
        
        # Debug
        if callback is not None: callback('downloading', len(data), len(chunks))
        
        if not quiet:
            print(f'[ BK ] Downloading chunk ({len(todo)} left)')
    
    # Check for missing chunks
    for i in range(len(chunks)):
        if not i in data.keys():
            
            # Debug
            if not quiet: print('Downloading missing', i)
            
            if callback is not None:
                callback('downloading missing', i, len(chunks))
            
            # Get chunk request
            mreq = [r for r in reqs if r.url.endswith(f'{i}.ts')][0]
            data[i] = dl_frag(mreq, timeout, session)
    
    return data

def bk_thread(raw: str,
              path: str,
              timeout: float = .05,
              callback: Callable = None,
              quiet: bool = False,
              **kwargs) -> None:
    '''
    Download each fragment using threads.
    Theorically faster.
    '''
    
    # Fetch the chunks
    chunks = bk_base_thread(raw = raw,
                            timeout = timeout,
                            callback = callback,
                            quiet = quiet)
    
    # Write them to the file
    with open(path, 'wb') as output:
        for key in sorted(chunks.keys()):
            
            if not quiet:
                print(f'\r[ BK ] Downloading: {key}/{len(chunks)}', end = '')
            
            if callback is not None:
                callback('writing', key, len(chunks))
            
            output.write(chunks[key])
        
        if not quiet: print()
    
    return
    
def bk_thread_ffmpeg(raw: str,
                     path: str,
                     timeout: float = .05,
                     callback: Callable = None,
                     quiet: bool = False,
                     **kwargs) -> None:
    '''
    Same as thread, but uses ffmpeg to concatenate.
    '''
    
    # Fetch the chunks
    chunks = bk_base_thread(raw = raw,
                            timeout = timeout,
                            callback = callback,
                            quiet = quiet)
    
    temp = 'temp/'
    track = []
    
    # Create temp folder if needed
    if not os.path.exists(temp): os.mkdir(temp)
    
    for index, chunk in chunks.items():
        
        frag_path = temp + str(index) + '.ts'
        
        with open(frag_path, 'wb') as frag:
            frag.write(chunk)
        
        if not quiet:
            print(f'\r[ BK ] Writing {index}/{len(chunks)}', end = '')
        
        if callback is not None:
            callback('writing', index, len(chunks))
        
        track += [f'file {index}.ts']
    
    with open(temp + 'track', 'w') as file:
        file.write('\n'.join(track))
    
    if not quiet: print('[ BK ] Concatenating')
    
    # Concatenate ts files with ffmpeg
    
    command = [
        FFMPEG, '-f', 'concat',
        '-safe', '0',
        '-i', temp + 'track',
        '-c', 'copy', path, '-y'
    ]
    
    def log(line: str) -> None:
        # Decide to execute the callback
        # for a line or not.
        
        if line.startswith('\nframe='):
            cur = line[8:].split()[0]
            
            if callback is not None:
                callback('concatenating', int(cur), None)
        
        if not quiet:
            print(line, end = '')
    
    assert utils.popen(command, log) == 0
    
    if not quiet: print('Deleting temp files')
    
    # Delete the temp files
    for file in os.listdir(temp):
        os.remove(temp + file)

def bk_safe(raw: str,
            path: str,
            timeout: float = 0,
            callback: Callable = None,
            quiet: bool = False,
            **kwargs) -> None:
    '''
    Download one segment at a time, and concatenate their
    bytes before writing to a file.
    '''
    
    fragments = re.findall(consts.re.fragments, raw)
    session = requests.Session()
    
    with open(path, 'wb') as output:
        for i, url in enumerate(fragments):
            
            if callback: callback('downloading', i, len(fragments))
            
            if not quiet:
                print('Downloading', i)
            
            output.write(
                session.get(url, headers = consts.headers).content
            )
            
            time.sleep(timeout)
    
    print()

# EOF
