#!/usr/bin/env python3.9
import argparse
import dataclasses
import datetime
import json
import logging
import re
import subprocess
import sys
import threading
import time
import typing
from pathlib import Path

import math
import rich.progress

logger = logging.getLogger(__name__)


@dataclasses.dataclass
class FfprobeResult:
    duration_sec: float
    duration_human: str
    codec: str
    fps: float
    size_bytes: int
    width: int
    height: int
    bitrate: int
    container: str
    ar: float
    sample_ar: float


def ffprobe(video: Path) -> FfprobeResult:
    logger.debug("executing ffprobe")
    proc = subprocess.run(
        args=[
            "ffprobe",
            "-v",
            "error",
            "-show_streams",
            "-show_format",
            "-print_format",
            "json",
            str(video),
        ],
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    try:
        proc.check_returncode()
    except subprocess.CalledProcessError:
        logger.error("failed to run ffprobe", exc_info=True)
        raise
    output: dict = json.loads(proc.stdout)
    video_stream: dict = [s for s in output.get("streams", []) if s.get("codec_type") == "video"][0]
    video_format: dict = output.get("format", {})
    width, height = video_stream.get("width"), video_stream.get("height")
    codec = video_stream.get("codec_name")
    duration = float(video_format.get("duration"))
    fps_long = eval(video_stream.get("r_frame_rate", ""))
    fps = float(f"{fps_long:.3f}")
    size = int(video_format.get("size", 0))
    duration_time = datetime.timedelta(seconds=math.ceil(duration))
    bitrate = int(video_stream.get("bit_rate", video_format.get("bit_rate", 0)))

    try:
        sample_w, sample_h = list(map(int, video_stream["sample_aspect_ratio"].split(":")))
        sample_ar = (sample_w / sample_h) or 1
    except:
        sample_ar = 1

    logger.debug(f"dimensions={width}x{height} duration={duration_time}, fps={fps}, size={size // 1_048_576}MB")
    # total_frames = video_stream.get("nb_frames")

    return FfprobeResult(
        duration_sec=duration,
        duration_human=str(duration_time),
        codec=codec,
        fps=fps,
        size_bytes=size,
        width=width,
        bitrate=bitrate,
        height=height,
        container=video.suffix.lstrip(".").lower(),
        ar=width / height,
        sample_ar=sample_ar,
        # total_frames=total_frames,
    )


def make_thumbnail_tile(
    video: Path,
    image_path: Path = None,
    columns: int = 3,
    interval_seconds: int = 60,
    tile_width: int = 540,
    skip_first_sec: int = 10,
    skip_if_exists: bool = False,
) -> Path:
    if not image_path:
        image_path = video.parent / f"{video.stem}.thumbnail.jpg"
    if image_path.is_file() and skip_if_exists:
        logger.debug("thumbnail already exists")
        return image_path

    info = ffprobe(video)

    min_frames = columns * 3
    sec_per_frame = min(interval_seconds, math.ceil(info.duration_sec / min_frames))
    rows = math.ceil(info.duration_sec // sec_per_frame / columns)
    tile = f"{columns}x{rows}"

    scaled_width, scaled_height = tile_width, math.ceil(tile_width / info.ar / info.sample_ar)
    scale = f"{scaled_width}:{scaled_height}"

    proc = subprocess.run(
        # fmt: off
        args=[
            "ffmpeg",
            "-v", "error",
            "-skip_frame", "nokey",
            "-ss", f"{skip_first_sec}",
            "-i", str(video.absolute()),
            "-vf", f"fps=1/{sec_per_frame},scale={scale},tile={tile}",
            "-frames", "1",
            "-y", str(image_path.absolute()),
        ],
        # fmt: on
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        timeout=2000,
    )
    proc.check_returncode()
    return image_path


@dataclasses.dataclass
class EncodeProgress:
    percent: float
    fps_avg: int
    eta: str
    current_size: int

    @property
    def encoded_mb(self) -> int:
        return math.ceil(self.current_size / 1_048_576)

    @property
    def estimated_mb(self) -> int:
        return math.ceil(self.encoded_mb / (self.percent / 100))


@dataclasses.dataclass
class UploadProgress:
    uploaded_mb: int
    speed: str
    percent: float


def noop(*args):
    pass


def upload_file(
    file_path: Path,
    destination: str,
    watch: bool = False,
    on_progress: typing.Callable[[UploadProgress], None] = noop,
    on_resync: typing.Callable[[], None] = noop,
):
    size = file_path.stat().st_size
    while True:
        proc = subprocess.Popen(
            ["rsync", "--bwlimit", "4500", "-rvPa", str(file_path), destination],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )

        def watch_progress(p: subprocess.Popen):
            #  5233864   0%    1.93MB/s    0:21:14
            #  12,450,960  30%    2.49MB/s    0:00:11
            pattern = re.compile(r"\b(?P<uploaded>[\d,]+)\s+(?P<percent>[\d]+)%\s+(?P<speed>\d+\.\d+[kmg]B/s)\b", re.I)
            for line in p.stdout:
                if m := pattern.search(line):
                    parsed = m.groupdict()
                    progress = UploadProgress(
                        uploaded_mb=int(parsed["uploaded"].replace(",", "")) // 1_048_576,
                        speed=parsed["speed"],
                        percent=float(parsed["percent"]),
                    )
                    try:
                        on_progress(progress)
                    except:
                        pass

        t = threading.Thread(target=watch_progress, args=(proc,))
        t.start()

        if proc.wait() != 0:
            code = proc.poll()
            err = proc.stderr.read()
            logger.warning(f"rsync failed with {code=}. error: {err}")
            return

        if not watch:
            return

        time.sleep(5)
        current_size = file_path.stat().st_size
        diff = current_size - size
        size = current_size

        if not diff:
            return

        logger.info(f"filesize changed {diff} bytes, re-syncing...")
        on_resync()


def encode_video_handbrake(
    video_path: Path,
    save_path: Path,
    quality: int = 30,
    from_time: typing.Optional[int] = None,
    duration: typing.Optional[int] = None,
    on_progress: typing.Callable[[EncodeProgress], None] = noop,
    extra_args: typing.List[str] = None,
    denoise: bool = False,
    is_10bit: bool = True,
):
    # fmt: off
    args = [
        'HandbrakeCLI',
        '--format', 'av_mp4',
        '--input', str(video_path),
        *(['--start-at', f'duration:{from_time}'] if from_time else []),
        *(['--stop-at', f'duration:{duration}'] if duration else []),
        '--output', str(save_path),
        '--optimize',
        '--encoder', *['vt_h265_10bit' if is_10bit else 'vt_h265'],
        '--quality', str(quality),
        '--vfr',
        '--aencoder', 'ac3',
        '--ab', '160',
        '--non-anamorphic',
        *(['--hqdn3d', 'light'] if denoise else []),
        *(extra_args or []),
        # '--json',
        '--verbose', '0'
    ]
    # fmt: on
    # subprocess.run(args, check=True)
    # return
    logger.debug("executing handbrake with args: %s", args)
    p = subprocess.Popen(args, stderr=subprocess.PIPE, stdout=subprocess.PIPE, text=True)

    last_progress: typing.Optional[EncodeProgress] = None

    def parse_progress(p: subprocess.Popen):
        # Encoding: task 1 of 1, 11.96 % (210.12 fps, avg 206.32 fps, ETA 00h06m14s)
        re_progress = re.compile(r"(?P<percent>[\d.]+) % \(.+ avg (?P<fps_avg>[\d.]+) fps, ETA (?P<eta>[^)]+)")
        for line in p.stdout:
            if m := re_progress.search(line.strip()):
                parsed = m.groupdict()
                percent = round(float(parsed["percent"]), 1)
                fps_avg = math.floor(float(parsed["fps_avg"]))
                eta = parsed["eta"]

                nonlocal last_progress
                last_progress = EncodeProgress(percent, fps_avg, eta, current_size=save_path.stat().st_size)

                try:
                    on_progress(last_progress)
                except:
                    pass

    t = threading.Thread(target=parse_progress, args=(p,))
    t.start()

    try:
        if p.wait() != 0:
            logger.error(f"handbrake failed: {p.stderr.read()}")
    except KeyboardInterrupt:
        p.kill()
        if last_progress:
            logger.info(f"cancelled at {last_progress.percent}%")
            logger.info(f"encoded file is {last_progress.encoded_mb}MB, estimated: {last_progress.estimated_mb}MB")
        raise


def encode_video_ffmpeg(
    video_path: Path,
    save_path: Path,
    quality: int = 30,
    from_time: typing.Optional[int] = None,
    duration: typing.Optional[int] = None,
    on_progress: typing.Callable[[EncodeProgress], None] = noop,
    extra_args: typing.List[str] = None,
    is_10bit: bool = True,
    **kwargs,
):
    # fmt: off
    args = [
        'ffmpeg',
        '-y',
        # '-v', 'quiet',
        '-progress', 'pipe:1',
        '-stats_period', '3',
        *(['-ss', str(from_time)] if from_time else []),
        *(['-t', str(duration)] if duration else []),
        '-i', str(video_path),
        '-c:v', 'hevc_videotoolbox',
        '-q:v', f'{quality}',
        '-profile:v', *['main10' if is_10bit else 'main'],
        '-map_metadata', '0',
        '-metadata', f'title={video_path.stem}',
        *(extra_args or []),
        str(save_path),
    ]
    # fmt: on

    logger.info(f"calling ffmpeg with {args=}")

    proc = subprocess.run(args, check=True)


def parse_args():
    arger = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    arger.add_argument("video_path", type=Path, help="path to video file")
    arger.add_argument("-e", "--encoder", choices=["handbrake", "ffmpeg"], default="handbrake", help="encoder engine")
    arger.add_argument("-q", "--quality", type=float, default=35, help="x265 quality factor")
    arger.add_argument("--output-dir", "-o", dest="output_dir", type=Path, help="Dir to save encoded files")
    arger.add_argument("--rsync", dest="upload_target", help="rsync encoded file to a host")
    arger.add_argument("--validate", action="store_true", default=False, help="perform validations before starting")
    arger.add_argument("--denoise", action="store_true", default=False, help="Enable denoise filter (Handbrake only)")
    arger.add_argument("--10bit", action="store_true", dest="is_10bit", help="Encode using 10-bit profile")
    arger.add_argument("--8bit", action="store_true", dest="is_10bit", help="Encode using 8-bit profile")
    arger.set_defaults(is_8bit=True)

    def parse_time(val: str) -> int:
        parts = val.split(":")
        if len(parts) == 1:
            return int(val)
        elif len(parts) == 2:
            return int(parts[0]) * 60 + int(parts[1])
        return int(parts[0]) * 60 + int(parts[1]) * 60 + int(parts[2])

    arger.add_argument(
        "--from",
        dest="from_time",
        default=0,
        type=parse_time,
        help="Start encoding from this time. Example 05:00 or 300",
    )
    arger.add_argument("--duration", type=parse_time, help="Stop encoding at this time. Example 07:00 or 420")

    if len(sys.argv[1:]) < 1:
        arger.print_help()
        exit(1)

    args, extra_args = arger.parse_known_args()
    if extra_args and extra_args[0] == "--":
        extra_args = extra_args[1:]

    return args, extra_args


def generate_filename(video_path: Path) -> str:
    probe = ffprobe(video_path)
    if 1900 <= probe.width <= 2000:
        hd = "1080p"
    elif 1200 <= probe.width <= 1400:
        hd = "720p"
    elif 3000 <= probe.width:
        hd = "4K"
    else:
        hd = None
    new_stem = re.sub(r"(\[\d+[pk]])", "", video_path.stem)
    new_stem = f"{new_stem.strip()} [{hd}, x265]"
    return new_stem


def main():
    args, extra_args = parse_args()

    video_path: Path = args.video_path
    video_path = video_path.expanduser().resolve()
    if not video_path.is_file():
        logger.error("No such file")
        exit(1)

    if args.validate and len(video_path.name) >= 128:
        logger.error("Filename is too long")
        exit(1)

    save_path = video_path.with_stem(generate_filename(video_path)).with_suffix(".mp4")
    if video_path == save_path:
        save_path = save_path.with_suffix(".reencoded" + video_path.suffix)

    output_dir: Path = args.output_dir
    if not output_dir:
        output_dir = save_path.parent / "_reenc"

    output_dir = output_dir.expanduser().resolve()
    output_dir.mkdir(exist_ok=True, parents=True)
    save_path = output_dir / save_path.name

    bar = rich.progress.Progress(refresh_per_second=2)
    task_encode = bar.add_task("encoding", visible=False)
    up_task = bar.add_task("uploading", visible=False, start=False)

    total_encoded_mb = 0

    def sync_later():
        logger.info("syncing in 10s")
        time.sleep(10)
        bar.update(up_task, visible=True)
        bar.start_task(up_task)
        upload_file(
            save_path,
            args.upload_target,
            watch=True,
            on_progress=lambda p: bar.update(
                up_task,
                description=f"upload: {p.speed}, {p.uploaded_mb:3}MB/{total_encoded_mb:3}MB",
                completed=p.uploaded_mb / total_encoded_mb * 100,
            ),
            on_resync=lambda: bar.reset(up_task, description="upload: re-syncing"),
        )

    t = threading.Thread(target=sync_later)
    if args.upload_target:
        t.start()

    encoder = encode_video_ffmpeg if args.encoder == "ffmpeg" else encode_video_handbrake

    with bar:

        def _on_encode_progress(p: EncodeProgress):
            nonlocal total_encoded_mb
            total_encoded_mb = p.encoded_mb
            bar.update(
                task_encode,
                completed=p.percent,
                description=f"encode: {p.fps_avg:3}fps, {p.encoded_mb:3}MB/{p.estimated_mb:3}MB",
            )

        try:
            bar.update(task_encode, visible=True)
            encoder(
                video_path,
                quality=args.quality,
                save_path=save_path,
                extra_args=extra_args,
                on_progress=_on_encode_progress,
                from_time=args.from_time,
                duration=args.duration,
                denoise=args.denoise,
                is_10bit=args.is_10bit,
            )
        except KeyboardInterrupt:
            logger.info("cancelled")
            exit(1)

        logger.info(f"finished encoding {video_path}")
        if args.upload_target:
            t.join()
            image_path = make_thumbnail_tile(save_path, image_path=save_path.with_suffix(".jpg"), skip_if_exists=True)
            upload_file(image_path, args.upload_target)


if __name__ == "__main__":
    logging.basicConfig(level=logging.DEBUG)
    main()