#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = ["httpx", "pillow", "click"]
# ///
import dataclasses
import logging
import os
import re
import struct
import subprocess
import sys
import urllib.parse
from functools import partial
from pathlib import Path
from typing import Callable

import click
import httpx
from PIL import Image

run_command = partial(
    subprocess.run,
    check=True,
)

# ---------------------------------------------------------------------------
# Image dimension reading (adapted from sort_images.py)
# ---------------------------------------------------------------------------

IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".gif"}


def read_dims_jpeg(image_path: Path) -> tuple[int, int]:
    with image_path.open("rb") as f:
        if f.read(2) != b"\xff\xd8":
            raise ValueError(f"{image_path} is not a valid JPEG file")
        f.seek(0)
        try:
            return Image.open(f).size
        except Exception:
            f.seek(0)
            while True:
                marker = f.read(1)
                if not marker or marker != b"\xff":
                    raise ValueError(f"Invalid JPEG format in {image_path}")
                marker_type = int.from_bytes(f.read(1), byteorder="big")
                length = int.from_bytes(f.read(2), byteorder="big") - 2
                is_sof = 0xC0 <= marker_type <= 0xCF and marker_type not in (
                    0xC4,
                    0xC8,
                    0xCC,
                )
                if is_sof:
                    f.seek(1, 1)
                    height = int.from_bytes(f.read(2), byteorder="big")
                    width = int.from_bytes(f.read(2), byteorder="big")
                    return width, height
                f.seek(length, 1)


def read_dims_png(image_path: Path) -> tuple[int, int]:
    with image_path.open("rb") as f:
        if f.read(8) != b"\x89PNG\r\n\x1a\n":
            raise ValueError(f"{image_path} is not a valid PNG file")
        try:
            f.seek(0)
            with Image.open(f) as img:
                return img.size
        except Exception:
            f.seek(12)
            if f.read(4) != b"IHDR":
                raise ValueError(f"IHDR chunk not found in {image_path}")
            dims = f.read(8)
            if len(dims) < 8:
                raise ValueError(f"Truncated IHDR in {image_path}")
            return struct.unpack(">II", dims)


def read_dims_webp(image_path: Path) -> tuple[int, int]:
    with image_path.open("rb") as f:
        header = f.read(12)
        if len(header) < 12 or header[:4] != b"RIFF" or header[8:12] != b"WEBP":
            raise ValueError(f"Not a valid WebP file: {image_path}")
        while True:
            chunk_header = f.read(8)
            if len(chunk_header) < 8:
                break
            tag, length = struct.unpack("<4sI", chunk_header)
            if tag == b"VP8X":
                data = f.read(10)
                return (int.from_bytes(data[4:7], "little") & 0xFFFFFF) + 1, (
                    int.from_bytes(data[7:10], "little") & 0xFFFFFF
                ) + 1
            elif tag == b"VP8L":
                data = f.read(5)
                if data[0] != 0x2F:
                    raise ValueError("Invalid VP8L signature")
                bits = struct.unpack("<I", data[1:5])[0]
                return (bits & 0x3FFF) + 1, ((bits >> 14) & 0x3FFF) + 1
            elif tag == b"VP8 ":
                data = f.read(10)
                if data[3:6] != b"\x9d\x01\x2a":
                    raise ValueError("Invalid VP8 sync code")
                w, h = struct.unpack("<HH", data[6:10])
                return w & 0x3FFF, h & 0x3FFF
            f.seek((length + 1) & ~1, 1)
    raise ValueError(f"Could not find dimension chunks in {image_path}")


def read_dims(image_path: Path) -> tuple[int, int]:
    match image_path.suffix.lower():
        case ".jpg" | ".jpeg":
            return read_dims_jpeg(image_path)
        case ".png":
            return read_dims_png(image_path)
        case ".webp":
            return read_dims_webp(image_path)
        case _:
            raise ValueError(f"Unsupported image format: {image_path.suffix}")


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def resolve_download_url_with_alldebrid(token: str, url: str) -> str:
    resp = httpx.get(
        "https://api.alldebrid.com/v4/link/unlock",
        params={"link": url},
        headers={"Authorization": f"Bearer {token}"},
        timeout=10,
    )
    resp.raise_for_status()
    data = resp.json()
    if data["status"] == "error":
        raise ValueError(f"{data['error']['code']}: {data['error']['message']}")
    return data["data"]["link"]


_CT_TO_EXT = {
    "application/zip": ".zip",
    "application/x-zip-compressed": ".zip",
    "application/x-rar-compressed": ".rar",
    "application/vnd.rar": ".rar",
    "application/x-7z-compressed": ".7z",
    "video/mp4": ".mp4",
    "video/quicktime": ".mov",
    "video/x-matroska": ".mkv",
}

_VIDEO_EXTS = {".mp4", ".mov", ".mkv"}


@dataclasses.dataclass
class UrlMeta:
    suggested_filename: str
    suggested_extension: str
    size_bytes: int


def _resolve_url(url: str) -> UrlMeta:
    """Send HEAD request; return filename, extension, and size from response headers."""
    try:
        resp = httpx.head(url, follow_redirects=True, timeout=15)
        cd = resp.headers.get("content-disposition", "")
        filename = None
        if "filename=" in cd:
            part = cd.split("filename=", 1)[1].strip().strip('"').strip("'")
            filename = urllib.parse.unquote(part.split(";")[0].strip())

        ct = resp.headers.get("content-type", "").split(";")[0].strip()
        ext = _CT_TO_EXT.get(ct) or Path(urllib.parse.urlparse(url).path).suffix

        if not filename:
            filename = urllib.parse.unquote(Path(urllib.parse.urlparse(url).path).name)
        if not Path(filename).suffix and ext:
            filename = filename + ext

        size = int(resp.headers.get("content-length", 0))
        return UrlMeta(
            suggested_filename=filename, suggested_extension=ext, size_bytes=size
        )
    except Exception:
        url_path = Path(urllib.parse.urlparse(url).path)
        name = urllib.parse.unquote(url_path.name)
        return UrlMeta(
            suggested_filename=name, suggested_extension=url_path.suffix, size_bytes=0
        )


def _name_from_url(url: str) -> str:
    parsed = urllib.parse.urlparse(url)
    segment = parsed.path.rstrip("/").split("/")[-1]
    return urllib.parse.unquote(segment) or parsed.netloc


def _has_images(d: Path) -> bool:
    return any(f.suffix.lower() in IMAGE_EXTS for f in d.rglob("*") if f.is_file())


def _needs_resize(d: Path) -> bool:
    images = [f for f in d.rglob("*") if f.is_file() and f.suffix.lower() in IMAGE_EXTS]
    if not images:
        return False
    mean_mb = sum(f.stat().st_size for f in images) / len(images)
    files_are_large = mean_mb >= 2 * 1_048_567
    images_are_large = any(max(read_dims(img)) > 5000 for img in images)

    return files_are_large or images_are_large


# ---------------------------------------------------------------------------
# Post-download: extract archives then resize
# ---------------------------------------------------------------------------


def _extract_and_process(output_dir: Path) -> None:
    archives = [
        f for pattern in ("*.zip", "*.rar", "*.7z") for f in output_dir.glob(pattern)
    ]
    for archive in archives:
        print(f"extracting {archive.name} …")
        run_command(["unar", "-o", str(output_dir), str(archive)])
        move_to_trash([archive])

    subdirs = [d for d in output_dir.iterdir() if d.is_dir() and d.name != "resized"]
    targets = subdirs if subdirs else [output_dir]

    to_process = [d for d in targets if _needs_resize(d)]
    if not to_process:
        print("no dirs need resizing")
    else:
        resized_dir = output_dir / "resized"
        resized_dir.mkdir(exist_ok=True)
        originals = [
            f
            for d in to_process
            for f in d.rglob("*")
            if f.is_file() and f.suffix.lower() in IMAGE_EXTS
        ]
        print(f"resizing {len(originals)} image(s) → {resized_dir}")
        for img in originals:
            run_command(
                [
                    "imgz",
                    "resize",
                    "-o",
                    str(resized_dir / f"{img.stem}.resized.jpg"),
                    str(img),
                ]
            )
        if originals:
            print(f"trashing {len(originals)} original(s)")
            move_to_trash(originals)
    rename_children(output_dir, flatten=True)
    remove_empty_dirs(output_dir)


def remove_empty_dirs(root_dir: Path) -> None:
    removed = [
        d
        for d in sorted(root_dir.rglob("*"), key=lambda p: len(p.parts), reverse=True)
        if d.is_dir() and not any(d.iterdir())
    ]
    for d in removed:
        d.rmdir()
    if removed:
        print(f"removed {len(removed)} empty dir(s) in {root_dir}")


def move_to_trash(file_paths: list[Path]) -> None:
    """Move files to Trash using trash CLI (single batched call)."""
    resolved: list[str] = []
    for file_path in file_paths:
        if not file_path.is_file():
            raise FileNotFoundError(f"File not found: {file_path}")
        resolved.append(str(file_path.resolve()))

    if resolved:
        run_command(["trash", *resolved])


def rename_children(
    source_dir: Path,
    flatten: bool,
    filter_fn: Callable[[Path], bool] | None = None,
) -> None:
    if not filter_fn:
        filter_fn = is_image

    def natural_sort_key(path: Path) -> tuple[str, ...]:
        s = path.stem
        return tuple(
            f"{int(p):010d}" if p.isnumeric() else p
            for p in re.findall(r"(\D+|\d+)", s)
        )

    assert source_dir.is_dir()

    if flatten:
        files = [f for f in source_dir.rglob("*") if f.is_file()]
    else:
        files = [f for f in source_dir.iterdir() if f.is_file()]

    files = [f for f in files if filter_fn(f)]

    files = sorted(files, key=natural_sort_key)

    for i, f in enumerate(files):
        ext = ".jpg" if f.suffix.lower() in {".jpeg", ".jpg"} else f.suffix
        target_dir = f.parent
        if flatten:
            target_dir = source_dir
        new_stem = source_dir.name
        target = target_dir / f"{new_stem}__{i:04d}{ext}"
        if target.is_file():
            logging.error(f"target already exists: {target}")
            continue
        f.rename(target)


def is_image(f: Path) -> bool:
    return f.is_file() and f.suffix.lower() in IMAGE_EXTS


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


_DEFAULT_CWD = Path("~/Downloads/_temp/_imagesets").expanduser()


@click.group()
def cli() -> None:
    pass


@cli.command("gallery")
@click.argument("urls", nargs=-1, required=True)
@click.option("--name", default=None)
@click.option("--cwd", type=Path, default=None)
def gallery_cmd(urls: tuple[str, ...], name: str | None, cwd: Path | None) -> None:
    urls = list(urls)
    resolved_name = name or _name_from_url(urls[0])
    output_dir = (cwd or Path(".")) / resolved_name

    if output_dir.exists() and _has_images(output_dir):
        print(f"skipping gallery-dl: {output_dir} already has images")
    else:
        print(f"gallery-dl → {output_dir}")
        args = ["uvx", "gallery-dl"]
        if name or cwd:
            args += ["-D", str(output_dir)]
        run_command(args + urls)

    _extract_and_process(output_dir)


@cli.command("filehost")
@click.argument("urls", nargs=-1, required=True)
@click.option("--name", default=None)
@click.option("--cwd", type=Path, default=_DEFAULT_CWD)
@click.option("--token", default=os.getenv("ALLDEBRID_TOKEN"))
def filehost_cmd(urls: tuple[str, ...], name: str | None, cwd: Path, token: str) -> None:
    resolved_name = name or _name_from_url(urls[0])
    output_dir = (cwd or Path(".")) / resolved_name
    output_dir.mkdir(parents=True, exist_ok=True)

    for url in urls:
        print(f"unrestricting {url}")
        dl_url = resolve_download_url_with_alldebrid(token, url)
        meta = _resolve_url(dl_url)
        if meta.suggested_extension in _VIDEO_EXTS:
            video_path = (cwd or Path(".")) / (resolved_name + meta.suggested_extension)
            print(f"video detected → {video_path}")
            if video_path.exists() and video_path.stat().st_size > 5 * 1024 * 1024:
                print(f"skipping: {video_path.name} already downloaded")
                continue
            run_command(["aria2c", "-o", str(video_path), dl_url])
            continue

        dest = output_dir / meta.suggested_filename
        if dest.exists() and dest.stat().st_size > 5 * 1024 * 1024:
            print(f"skipping: {dest.name} already downloaded")
            continue
        print(f"downloading {meta.suggested_filename} → {output_dir}")
        run_command(["aria2c", f"--dir={output_dir}", dl_url])

    _extract_and_process(output_dir)


@cli.command("download")
@click.argument("urls", nargs=-1, required=True)
@click.option("--name", default=None)
@click.option("--cwd", type=Path, default=_DEFAULT_CWD)
@click.option("--token", default=os.getenv("ALLDEBRID_TOKEN"))
@click.pass_context
def download_cmd(
    ctx: click.Context, urls: tuple[str, ...], name: str | None, cwd: Path, token: str
) -> None:
    filehost_urls = [u for u in urls if "rg.to" in u or "rapidgator.net" in u]
    gallery_urls = [u for u in urls if u not in set(filehost_urls)]
    if filehost_urls:
        ctx.invoke(filehost_cmd, urls=tuple(filehost_urls), name=name, cwd=cwd, token=token)
    if gallery_urls:
        ctx.invoke(gallery_cmd, urls=tuple(gallery_urls), name=name, cwd=cwd)


@cli.command("process")
@click.argument("dirs", nargs=-1, required=True)
def process_cmd(dirs: tuple[str, ...]) -> None:
    for d in dirs:
        p = Path(d)
        if not p.exists():
            print(f"skipping {p}: does not exist", file=sys.stderr)
            continue
        if _needs_resize(p):
            resized_dir = p / "resized"
            resized_dir.mkdir(exist_ok=True)
            originals = [f for f in p.rglob("*") if is_image(f)]
            print(f"resizing {len(originals)} image(s) in {p} → {resized_dir}")
            for img in originals:
                run_command(
                    [
                        "imgz",
                        "resize",
                        "-o",
                        str(resized_dir / f"{img.stem}.resized.jpg"),
                        str(img),
                    ]
                )
            if originals:
                print(f"trashing {len(originals)} original(s)")
                move_to_trash(originals)
            remove_empty_dirs(p)
        else:
            print(f"skipping {p}: already within size limits")


if __name__ == "__main__":
    cli()