#!/usr/bin/env python3
import argparse
import os
import re
import shutil
import subprocess
import sys
from pathlib import Path

SHOWS_ROOT = Path(__file__).resolve().parent
DOWNLOADS = SHOWS_ROOT / ".downloads"
SEED_RATIO = 2.0
SEED_TIME_MIN = 1440
PORT = 6881
VIDEO_EXTS = {'.mkv', '.mp4', '.avi', '.mov', '.m4v', '.ts', '.webm'}

def extract_ep(name):
    stem = Path(name).stem
    stem = re.sub(r'^\[[^\]]+\]\s*', '', stem)
    stem = re.sub(r'\([^)]*\)', '', stem)
    stem = re.sub(r'\[[^\]]*\]', '', stem)
    if re.search(r'\b(NCOP|NCED|OP|ED|Menu|Extra|Preview|PV|SP|Special)\b', stem, re.I):
        return None
    m = re.search(r'(?:^|\s|-)\s*(\d{1,3})(?:v\d)?(?=\s|$|\.)', stem)
    return int(m.group(1)) if m else None

def organize(stage, name, year, season):
    show_dir = SHOWS_ROOT / f"{name} ({year})"
    season_dir = show_dir / f"Season {int(season):02d}"
    specials_dir = show_dir / "Specials"
    season_dir.mkdir(parents=True, exist_ok=True)
    for f in Path(stage).rglob('*'):
        if not f.is_file() or f.suffix.lower() not in VIDEO_EXTS:
            continue
        ep = extract_ep(f.name)
        if ep is None:
            specials_dir.mkdir(parents=True, exist_ok=True)
            dst = specials_dir / f.name
        else:
            dst = season_dir / f"{name} - S{int(season):02d}E{ep:02d}{f.suffix.lower()}"
        if dst.exists():
            continue
        try:
            os.link(f, dst)
        except OSError:
            shutil.copy2(f, dst)
        print(f"[organize] {dst.relative_to(SHOWS_ROOT)}", flush=True)

def cmd_organize(a):
    organize(a.stage, a.name, a.year, a.season)

def cmd_download(a):
    torrent = Path(a.torrent).expanduser().resolve()
    if not torrent.exists():
        sys.exit(f"not found: {torrent}")
    stage = DOWNLOADS / torrent.stem
    stage.mkdir(parents=True, exist_ok=True)
    hook = stage / "_hook.sh"
    hook.write_text(
        f'#!/bin/bash\n'
        f'python3 "{Path(__file__).resolve()}" organize '
        f'--stage "{stage}" --name "{a.name}" --year "{a.year}" --season {a.season}\n'
    )
    hook.chmod(0o755)
    cmd = [
        "aria2c", str(torrent),
        f"--dir={stage}",
        f"--listen-port={PORT}", f"--dht-listen-port={PORT}",
        "--enable-dht=true", "--enable-peer-exchange=true", "--bt-enable-lpd=true",
        "--follow-torrent=mem", "--check-integrity=true",
        "--summary-interval=30", "--console-log-level=warn",
        f"--on-bt-download-complete={hook}",
    ]
    if a.no_seed:
        cmd.append("--seed-time=0")
    else:
        cmd += [f"--seed-ratio={SEED_RATIO}", f"--seed-time={SEED_TIME_MIN}"]
    subprocess.run(cmd)

def main():
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)
    d = sub.add_parser("download")
    d.add_argument("torrent")
    d.add_argument("--name", required=True)
    d.add_argument("--year", required=True)
    d.add_argument("--season", type=int, default=1)
    d.add_argument("--no-seed", action="store_true")
    d.set_defaults(func=cmd_download)
    o = sub.add_parser("organize")
    o.add_argument("--stage", required=True)
    o.add_argument("--name", required=True)
    o.add_argument("--year", required=True)
    o.add_argument("--season", type=int, default=1)
    o.set_defaults(func=cmd_organize)
    a = ap.parse_args()
    a.func(a)

if __name__ == "__main__":
    main()