#!/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 shutil
import struct
import subprocess
import urllib.parse
import zipfile
from functools import partial
from pathlib import Path
from typing import Callable

import click
import httpx
from PIL import Image, ImageOps

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 and not f.is_relative_to(d / "resized")]
    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 = False
    for image in images:
        try:
            if max(read_dims(image)) > 5000:
                images_are_large = True
                break
        except Exception as error:
            logging.warning(f"could not read dimensions of {image}; trying resizer: {error}")
            images_are_large = True
            break

    return files_are_large or images_are_large


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


ImageProcessor = Callable[..., None]


def detect_resizer() -> ImageProcessor:
    vips_binary = shutil.which(cmd="vips")
    if vips_binary:
        return partial(resize_with_vips, vips_binary=vips_binary)
    return resize_with_pillow


def resize_with_vips(image_path: Path, output_path: Path, vips_binary: str) -> None:
    try:
        subprocess.run(
            args=[
                vips_binary,
                "thumbnail",
                str(image_path),
                f"{output_path}[Q=75]",
                "5000",
                "--height",
                "5000",
                "--size",
                "down",
            ],
            check=True,
            capture_output=True,
        )
    except (OSError, subprocess.CalledProcessError) as error:
        logging.warning(f"vips failed for {image_path}; using Pillow: {error}")
        resize_with_pillow(image_path=image_path, output_path=output_path)


def resize_with_pillow(image_path: Path, output_path: Path) -> None:
    with Image.open(fp=image_path) as image:
        with ImageOps.exif_transpose(image=image) as oriented:
            oriented.thumbnail(size=(5000, 5000), resample=Image.Resampling.LANCZOS)
            with oriented.convert(mode="RGB") as resized:
                resized.save(fp=output_path, format="JPEG", quality=75)


def process_imageset(path: Path, processor: ImageProcessor) -> None:
    if not path.exists():
        logging.warning(f"skipping {path}: does not exist")
        return
    if path.is_file() and path.suffix.lower() in {".zip", ".rar", ".7z"}:
        output_dir = path.with_suffix("")
        archives = [path]
    elif path.is_dir():
        output_dir = path
        archives = [f for f in path.iterdir() if f.is_file() and f.suffix.lower() in {".zip", ".rar", ".7z"}]
    else:
        logging.warning(f"skipping {path}: expected a directory or archive")
        return
    if output_dir.exists() and not output_dir.is_dir():
        logging.warning(f"skipping {path}: {output_dir} is not a directory")
        return

    for archive in archives:
        logging.info(f"extracting {archive} → {output_dir}")
        if archive.suffix.lower() == ".zip":
            try:
                with zipfile.ZipFile(file=archive) as zipped:
                    zipped.extractall(path=output_dir)
            except zipfile.BadZipFile:
                logging.warning(f"skipping {archive}: invalid ZIP file")
                return
        else:
            run_command(args=["unar", "--output-directory", str(output_dir), str(archive)])

    output_dir.mkdir(parents=True, exist_ok=True)
    failed_images: set[Path] = set()
    if _needs_resize(output_dir):
        resized_dir = output_dir / "resized"
        resized_dir.mkdir(exist_ok=True)
        originals = [f for f in output_dir.rglob("*") if is_image(f) and not f.is_relative_to(resized_dir)]
        logging.info(f"resizing {len(originals)} image(s) → {resized_dir}")
        resized_originals: list[Path] = []
        for img in originals:
            output_path = resized_dir / f"{img.stem}.resized.jpg"
            if output_path.is_file() and output_path.stat().st_size > 0:
                logging.info(f"skipping {img}: {output_path} already resized")
                resized_originals.append(img)
                continue
            try:
                processor(image_path=img, output_path=output_path)
            except Exception as error:
                logging.warning(f"skipping {img}: resizing failed: {error}")
                failed_images.add(img)
                output_path.unlink(missing_ok=True)
                continue
            resized_originals.append(img)
        if resized_originals:
            logging.info(f"trashing {len(resized_originals)} original(s)")
            move_to_trash(file_paths=resized_originals)
    else:
        logging.info(f"skipping resizing {output_dir}: already within size limits")

    rename_children(
        source_dir=output_dir,
        flatten=True,
        filter_fn=lambda image: is_image(image) and image not in failed_images,
    )
    remove_empty_dirs(root_dir=output_dir)
    if archives and not failed_images:
        move_to_trash(file_paths=archives)


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)

    process_imageset(path=output_dir, processor=detect_resizer())


@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])

    process_imageset(path=output_dir, processor=detect_resizer())

def is_filehost(url: str) -> bool:
    kws = ["rg.to", "rapidgator.net", "trbt.cc", "katfile"]
    return any(kw in url for kw in kws)

@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 is_filehost(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("paths", nargs=-1, required=True, type=click.Path(path_type=Path))
def process_cmd(paths: tuple[Path, ...]) -> None:
    """Process directories and archives, extracting archives into sibling directories."""
    processor = detect_resizer()
    for path in paths:
        process_imageset(path=path, processor=processor)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
    cli()