from logging import getLogger
from pathlib import Path
from time import time
from typing import Any

import ffmpeg
from rich.prompt import Prompt

from utils.log import configure_logging

PATH = Path(__file__).parent
INPUT_DIR = PATH / 'input'
OUTPUT_DIR = PATH / 'output'

OUTPUT_ARGS: dict[str, Any] = {
    'row-mt': 1,
    'loglevel': 'error',
}


logger = getLogger(__name__)


def compress(input_file: Path, output_file: Path, sticker_width: int, max_size: int):
    file_input = ffmpeg.input(input_file)
    file_output = OUTPUT_DIR / input_file.name.replace(input_file.suffix, '.webm')

    stream = file_input.fps(fps=25, round='up')
    stream = stream.scale(w=sticker_width, h=-1, force_original_aspect_ratio='decrease')
    stream = stream.loop(loop=0)
    stream = stream.crop(
        out_w=f'min(min(iw,ih),{sticker_width})', out_h=f'min(min(iw,ih),{sticker_width})'
    )
    stream = stream.pad(color='0x00000000')

    if input_file.suffix in ('.gif', '.mp4'):
        duration = float(ffmpeg.probe(input_file)['streams'][0]['duration'])
        speed_factor = 1

        if duration > 3:
            speed_factor = duration / 3

        stream = stream.setpts(expr=f'PTS*{speed_factor}')

    output = ffmpeg.output(
        stream,
        filename=file_output,
        vcodec='libvpx-vp9',
        pix_fmt='yuva420p',
        extra_options=OUTPUT_ARGS,
    )

    try:
        output.run(overwrite_output=True)
    except Exception as e:
        logger.error(f"Error during first encoding: {e}")
        exit(1)

    file_size = file_output.stat().st_size

    if file_size / 1024 > max_size:
        initial_duration = float(ffmpeg.probe(input_file)['format']['duration'])
        bitrate = int(max_size * 8 / initial_duration * .65)

        temp_output = file_output.with_suffix('.tmp.webm')
        stream = ffmpeg.input(output_file, c='libvpx-vp9')

        output = ffmpeg.output(
            stream,
            filename=temp_output,
            vcodec='libvpx-vp9',
            pix_fmt='yuva420p',
            b=f'{bitrate}k',
            extra_options=OUTPUT_ARGS,
        )

        output.run(overwrite_output=True)

        file_output.unlink()
        temp_output.rename(output_file)


def main():
    if not INPUT_DIR.exists():
        INPUT_DIR.mkdir(parents=True)

    if not OUTPUT_DIR.exists():
        OUTPUT_DIR.mkdir(parents=True)

    for file in OUTPUT_DIR.iterdir():
        file.unlink()

    sticker_type = Prompt.ask(
        'Do you want to create stickers or emoji or icon?', choices=['s', 'e', 'i']
    )

    match sticker_type:
        case 's':
            width = 512
            size = 256
        case 'e':
            width = 100
            size = 256
        case 'i':
            width = 100
            size = 32
        case _:
            width = 512
            size = 256

    started_at = time()

    for file in INPUT_DIR.iterdir():
        logger.info(f'Processing {file.name}...')

        if file.suffix not in ('.png', '.jpg', '.webp', '.gif', '.mp4'):
            logger.warning(f'Unsupported file type: {file.suffix}')
            continue

        file_path = INPUT_DIR / file.name
        file_output = OUTPUT_DIR / file.name.replace(file.suffix, '.webm')

        compress(file_path, file_output, width, size)

    logger.info(f'Finished in {round(time() - started_at, 2)} seconds')


if __name__ == '__main__':
    configure_logging()
    main()