#!/usr/bin/env python3.9
import argparse
import json
import logging
import os
import re
import select
import shutil
import sys
import time
from turtle import back
import typing
from pathlib import Path
from typing import Optional

import ffmpeg
from ffmpeg_strip import clean_video

logging.basicConfig(level=logging.INFO, format=f"%(asctime)s {logging.BASIC_FORMAT}")


def force_import(module: str):
    import importlib
    import subprocess
    import sys

    try:
        return importlib.import_module(module)
    except ModuleNotFoundError:
        subprocess.run([sys.executable, "-m", "pip", "install", module])
        importlib.invalidate_caches()
        return importlib.import_module(module)


try:
    import torf
except ImportError:
    torf = force_import("torf")

try:
    import httpx
except ImportError:
    httpx = force_import("httpx")

try:
    import inquirer
except ImportError:
    inquirer = force_import("inquirer")

TORRENT_TRACKER_URL = os.getenv(
    "TORRENT_TRACKER_URL", "http://tracker.empornium.sx:2710/tegqucis10uanp672qh6nhn393xlkncs/announce"
)
TORRENT_CREATED_BY = os.getenv("TORRENT_CREATED_BY", "zzzp")
TORRENT_DIR = Path(os.getenv("TORRENT_DIR", "/mnt/box/files/_torrents/_new2/"))

VIDEO_EXTENSIONS = {".mp4", ".mpg", ".mkv"}


def make_torrent(source_dir: Path, tracker_url: str = TORRENT_TRACKER_URL) -> Path:
    torrent_path = source_dir / f"{source_dir.name}.torrent"
    if torrent_path.is_file():
        logging.info("Torrent file is already created")
        return torrent_path

    t = torf.Torrent(
        path=str(source_dir.resolve()),
        name=source_dir.resolve().name,
        trackers=[tracker_url],
        private=True,
        created_by=TORRENT_CREATED_BY,
        exclude_globs=["*.txt", "post.txt", "*.torrent", "*.gif"],
    )
    t.generate()
    t.write(torrent_path, overwrite=True)

    return torrent_path


def human_size(size: int) -> str:
    suffix = "B"
    for unit in ["", "K", "M", "G", "T", "P", "E", "Z"]:
        if abs(size) < 1024.0:
            return "%3.2f%s%s" % (size, unit, suffix)
        size /= 1024.0
    return "%.2f%s%s" % (size, "Yi", suffix)


def find_video(root_dir: Path) -> Optional[Path]:
    for ext in VIDEO_EXTENSIONS:
        for f in root_dir.rglob(f"*{ext}"):
            return f
    return None


def pick_name(video_path: Path) -> str:
    try:
        if parse_filename(video_path.stem):
            return video_path.stem
    except ValueError:
        pass

    candidates = set()

    parent = video_path.parent
    while parent.exists():
        try:
            if parse_filename(parent.name):
                candidates.add(parent.name)
        except ValueError:
            pass
        if parent.parent == parent:
            break
        parent = parent.parent

    if len(candidates) == 1:
        return candidates.pop()
    elif len(candidates) > 1:
        questions = [
            inquirer.List(
                "best_name",
                message="Pick the best filename",
                choices=[*candidates, "<custom>"],
            ),
        ]

        answers = inquirer.prompt(questions) or {}
        best_name = answers.get("best_name")
        if best_name != "<custom>":
            return best_name

    while True:
        try:
            answers = inquirer.prompt([inquirer.Editor("best_name", message="New name", default=video_path.stem)])
            if not answers:
                raise Exception("Cancelled")
            best_name = answers["best_name"].splitlines(keepends=False)[0].strip()
            parse_filename(best_name)
            return best_name
        except ValueError:
            continue


def move_video(video_path: Path, new_name: typing.Optional[str], link: bool = False) -> Path:
    """
    Move video to the torrent directory and returns the video path
    """
    if new_name:
        best_name = Path(new_name).name
    else:
        best_name = pick_name(video_path)
    assert not best_name.startswith('.')
    logging.info("using name: %s", best_name)

    best_name = best_name.removesuffix('.nometadata')
    if Path(best_name).suffix in VIDEO_EXTENSIONS:
        best_name = Path(best_name).stem

    renamed_path = video_path.with_stem(best_name)
    if link:
        video_path = video_path.link_to(renamed_path)
    else:
        video_path = video_path.rename(renamed_path)
    video_path = renamed_path

    new_loc = TORRENT_DIR / video_path.stem
    new_loc.mkdir(parents=True, exist_ok=True)
    target_path = new_loc / video_path.name
    video_path.rename(target_path)

    for f in video_path.parent.glob("*.jpg"):
        if f.stem.startswith(video_path.stem):
            f.rename(new_loc / f.name)
            break

    return target_path


def upload_image(image_path: Path) -> str:
    session = httpx.Client()
    res = session.get("https://jerking.empornium.ph/?agree-consent", follow_redirects=True)
    try:
        auth_token = re.search(r'name="auth_token" value="([^"]+)"', res.text).group(1)
    except AttributeError:
        auth_token = re.search(r'auth_token = "([^"]+)"', res.text).group(1)

    with image_path.open("rb") as f:
        res = session.post(
            "https://jerking.empornium.ph/json",
            headers={"accept": "application/json"},
            data={
                "thumb_width": "160",
                "thumb_height": "160",
                "thumb_crop": "false",
                "medium_width": "500",
                "medium_crop": "false",
                "type": "file",
                "action": "upload",
                "timestamp": str(round(time.time() * 1000)),
                "auth_token": auth_token,
                "nsfw": "0",
            },
            files={
                "source": ("thumbs.jpg", f),
            },
        )

    data = res.json()
    image_url = data["image"]["url"]
    thumbnail_url = data["image"]["display_url"]

    return image_url


ParsedFilename = typing.TypedDict(
    'ParsedFilename', {'actors': list[str], 'studio': str, 'date': str, 'title': str, 'tags': list[str]}
)


def parse_filename(filename: str) -> ParsedFilename:
    parsed = {
        "tags": [],
    }
    filename = re.sub(r"\s*\[\d+[^]]+]", "", filename)
    if m := re.search(r"(?P<actors>.+)\s+-+\s+@(?P<studio>\S+)\s+-+\s+(?P<title>.+)\s+-+\s+(?P<date>[\d-]+)", filename):
        parsed.update(m.groupdict())
    elif m := re.search(r"(?P<actors>.+)\s+-+\s+@(?P<studio>\S+)\s+-+\s+(?P<date>[\d-]+)", filename):
        parsed.update(m.groupdict())
    elif m := re.search(r"(?P<actors>.+)\s+-+\s+(?P<title>.+)\s+-+\s+(?P<date>[\d-]+)", filename):
        parsed.update(m.groupdict())
    elif m := re.search(r"(?P<actors>.+)\s+-+\s+(?P<title>\D.+)", filename):
        parsed.update(m.groupdict())
    else:
        raise ValueError("Unknown filename format")

    if "actors" in parsed:
        names = re.split(r"\s*,\s*", parsed["actors"])
        with_aliases = []
        for it in names:
            with_aliases.extend(re.split(r"\s+aka\s+", it))
        parsed["actors"] = names
        parsed["tags"].extend([it.lower().replace(" ", ".") for it in with_aliases])
    if "studio" in parsed:
        parsed["tags"].append(parsed["studio"].lower() + ".com")
    if "date" in parsed:
        year, ym = parsed["date"][:4], parsed["date"][:7].replace("-", ".")
        parsed["tags"].append(year)
        parsed["tags"].append(ym)

    return parsed


def generate_title(parsed: ParsedFilename) -> str:
    parts = []
    if actors := parsed.get('actors'):
        parts.append(', '.join(actors))
    if studio := parsed.get('studio'):
        parts.append(f'@{studio}')
    if title := parsed.get('title'):
        parts.append(title)
    if date := parsed.get('date'):
        parts.append(date)
    return ' -- '.join(parts)


def generate_post_bbcode(video_path: Path, thumbnails_path: Path) -> str:
    try:
        image_url = upload_image(thumbnails_path)
        image_bbcode = f"[img]{image_url}[/img]"
    except:
        logging.exception("Failed to upload image")
        image_bbcode = ""

    parsed = parse_filename(video_path.stem)
    probe = ffmpeg.ffprobe(video_path)

    if 1900 <= probe.width <= 2200:
        hd = "1080p"
    elif 1200 <= probe.width <= 1400:
        hd = "720p"
    elif 3000 <= probe.width:
        hd = "4K"
    else:
        hd = None

    if hd:
        parsed["tags"].append(hd.lower())
    if probe.codec == "hevc":
        parsed["tags"].extend(["x265", "x265.reencode", "hevc.x265"])

    table = {
        "Duration": probe.duration_human,
        "Format": video_path.suffix.strip("."),
        "Filesize": human_size(video_path.stat().st_size),
        "Resolution": f"{probe.width}x{probe.height}",
        "Codec": probe.codec,
        "Bit rate": f"{probe.bitrate // 1000} kbit/s",
        "FPS": probe.fps,
    }
    release_date = parsed.get("date", "")

    if release_date:
        table = {
            "Release Date": release_date,
            **table,
        }

    lines = [
        "[table=nball]",
        *(f"[tr][th=20]{k}[/th][td]{v}[/td][/tr]" for k, v in table.items()),
        "[/table]",
    ]
    table_bbcode = "\n".join(lines)

    m_performers = re.search(r"(.+?) -+\s", video_path.stem)
    performers = m_performers.group(1) if m_performers else ""

    title = generate_title(parsed)

    metadata = {
        **parsed,
        "is_hevc": probe.codec == "hevc",
        "hd": hd,
    }
    metadata_json = json.dumps(metadata)

    return f"""
{metadata_json}
---
[b]{title}[/b]

[cast]
{performers}

[details]

[info]
{table_bbcode}

[screens]
{image_bbcode}
""".strip()


def add_torrent(torrent_path: Path):
    import time
    from functools import partial

    make_id = partial(time.time_ns)

    # log in
    logging.debug("Connecting to Deluge")
    session = httpx.Client(base_url="https://t.zzzp.win/", timeout=30)
    res = session.post("/json", json={"method": "auth.login", "params": ["xAsametk50"], "id": make_id()})
    res.raise_for_status()

    # find first available host
    res = session.post("/json", json={"method": "web.get_hosts", "params": [], "id": make_id()})
    res.raise_for_status()
    host_id = res.json()["result"][0][0]

    # connect to a host
    res = session.post("/json", json={"method": "web.connect", "params": [host_id], "id": make_id()})
    res.raise_for_status()

    # upload torrent
    logging.debug("Uploading torrent file")
    with torrent_path.open("rb") as f:
        res = session.post("/upload", files={"file": f})
    res.raise_for_status()
    remote_path: str = res.json()["files"][0]

    # add torrent
    logging.debug("Adding torrent file")
    res = session.post(
        "/json",
        json={
            "method": "web.add_torrents",
            "params": [
                [
                    {
                        "path": remote_path,
                        "options": {
                            "file_priorities": [1],
                            "add_paused": True,
                            "sequential_download": False,
                            "pre_allocate_storage": False,
                            "download_location": "/dl/_new2",
                            "move_completed": False,
                            "move_completed_path": "/root/Downloads",
                            "prioritize_first_last_pieces": True,
                            "seed_mode": True,
                            "super_seeding": False,
                        },
                    }
                ]
            ],
            "id": make_id(),
        },
    )

    res.raise_for_status()
    torrent_id = res.json()["result"][0][1]

    # force recheck
    # logging.info("Triggering a forced recheck")
    # time.sleep(5)
    # try:
    #     _ = session.post(
    #         "/json", json={"method": "core.force_recheck", "params": [[torrent_id]], "id": make_id()}, timeout=0.1
    #     )
    # except httpx.HTTPError:
    #     pass

    logging.debug("Torrent file has been added successfully.")


def parse_args(argv: list[str]):
    arger = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    arger.add_argument("dir_or_video", type=Path, help="Directory or video file to add")
    arger.add_argument("--name", "-n", help="Torrent name")
    arger.add_argument("--no-metadata", action='store_true', dest='remove_metadata', help="Remove all metadata")
    arger.add_argument("--post-only", action="store_true", help="Generate only the post bbcode")
    arger.add_argument("--link", action="store_true", help="Link video file instead of moving")
    arger.add_argument("--thumb-path", dest='thumbnail_path', type=Path, help="Path to thumbnail sheet")
    if not argv:
        arger.print_help()
        sys.exit(0)
    return arger.parse_args(argv)


def validate_name(base_name: str):
    if len(base_name) > (128 - len(".mp4")):
        raise ValueError(f"Filename is too long: {base_name} ({len(base_name)} chars)")
    if "?" in base_name:
        raise ValueError(f"Filename contains invalid characters: {base_name}")


def main():
    args = parse_args(sys.argv[1:])
    dir_or_video: Path = args.dir_or_video.expanduser().resolve()
    if dir_or_video.is_dir():
        dir_path = dir_or_video
        if not dir_path.is_relative_to(TORRENT_DIR):
            target_path = TORRENT_DIR / dir_path.name
            logging.info(f"Linking to {target_path}")
            shutil.copytree(src=dir_path, dst=target_path, copy_function=os.link, symlinks=False)
            dir_path = target_path
        torrent_path = make_torrent(dir_path)
        logging.info(f"Saved torrent at {torrent_path}")

        logging.info("Adding torrent file to Deluge")
        add_torrent(torrent_path)
        return

    if not dir_or_video.is_file():
        logging.error("Invalid file or directory")
        return

    video_path = dir_or_video
    if video_path.suffix not in VIDEO_EXTENSIONS:
        logging.error('not a video')
        return

    if args.remove_metadata:
        logging.info('Removing video metadata')
        save_path = video_path.with_stem(f'{video_path.stem}.nometadata')
        video_path = clean_video(video_path=video_path, save_path=save_path)

    video_path = move_video(video_path, new_name=args.name, link=args.link)
    dir_path = video_path.parent

    if p := args.thumbnail_path:
        thumbnail_path = video_path.with_suffix('.jpg')
        p.link_to(thumbnail_path)
    else:
        logging.info("Creating thumbnail tile")
        old_thumb_path = video_path.with_suffix(".thumbnail.jpg")
        thumbnail_path = video_path.with_suffix(".jpg")
        if old_thumb_path.is_file():
            old_thumb_path.rename(thumbnail_path)
        image_jpeg = ffmpeg.make_thumbnail_tile(video_path)
        thumbnail_path.write_bytes(image_jpeg)
        logging.info(f"Saved thumbnails at {thumbnail_path}")

    post_bbcode = generate_post_bbcode(video_path, thumbnail_path)
    video_path.with_name("post.txt").write_text(post_bbcode)
    print(post_bbcode)

    if args.post_only:
        return

    total_files = len(list(dir_path.rglob("*")))
    logging.info(f"Creating torrent file from {dir_path}. Total files: {total_files}")
    torrent_path = make_torrent(dir_path)
    logging.info(f"Saved torrent at {torrent_path}")

    logging.info("Adding torrent file to Deluge")
    add_torrent(torrent_path)


if __name__ == "__main__":
    main()