#!/usr/bin/env -S uv run --script
# /// script
# dependencies = []
# ///

import argparse
import logging
import subprocess
from threading import Thread
from typing import Callable
from pathlib import Path

GIGAPIXEL_EXE = Path(
    "/Applications/Topaz Gigapixel AI.app/Contents/MacOS/Topaz Gigapixel AI"
)
MAX_LENGTH = 4000


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Upscale images using Topaz Gigapixel AI CLI.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "images",
        nargs="+",
        type=Path,
        help="One or more image files to upscale.",
    )
    parser.add_argument(
        "--format",
        choices=("jpg", "png"),
        default="jpg",
        help="Output image format.",
    )
    parser.add_argument(
        "--max-length",
        type=int,
        default=4000,
        help="Maximum length of the longest side of the output image.",
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Regenerate images even if upscaled outputs already exist.",
    )
    return parser.parse_args()


def filter_out_upscaled(
    images: list[Path], force: bool, output_format: str
) -> list[Path]:
    if force:
        return images

    filtered: list[Path] = []
    for img_path in images:
        if img_path.stem.endswith("-upscaled"):
            continue
        output_dir = img_path.parent
        output_stem = f"{img_path.stem}-upscaled"
        upscaled_candidates = [output_dir / f"{output_stem}.{output_format}"]
        if output_format == "jpg":
            upscaled_candidates.append(output_dir / f"{output_stem}.jpeg")
        if any(candidate.exists() for candidate in upscaled_candidates):
            continue
        filtered.append(img_path)

    return filtered


def upscale_images(
    images: list[Path],
    output_dir: Path,
    output_format: str,
    on_progress: Callable[[Path], None],
) -> None:
    cmd = [
        str(GIGAPIXEL_EXE),
        "--cli",
        "--verbose",
        "--parallel",
        "2",
        "--input",
        *[str(img) for img in images],
        "--output",
        str(output_dir),
        "--height",
        str(MAX_LENGTH),
        "--image-format",
        output_format,
        "--suffix",
        "-upscaled",
    ]

    if output_format == "jpg":
        cmd.extend(["--jpeg-quality", "75"])

    def stream_stdout(process: subprocess.Popen[str]) -> None:
        if process.stdout is None:
            return
        for line in process.stdout:
            if "Saved file:" not in line:
                continue
            parts = line.split('"')
            if len(parts) < 3:
                continue
            saved_path = parts[-2].strip()
            if saved_path:
                on_progress(Path(saved_path))

    process = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        bufsize=1,
    )
    stdout_thread = Thread(target=stream_stdout, args=(process,), daemon=True)
    stdout_thread.start()
    stderr = process.stderr.read() if process.stderr is not None else ""
    return_code = process.wait()
    stdout_thread.join()
    if return_code != 0:
        raise subprocess.CalledProcessError(return_code, cmd, stderr=stderr)


def main() -> None:
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s"
    )
    args = parse_args()
    global MAX_LENGTH
    MAX_LENGTH = args.max_length

    if not GIGAPIXEL_EXE.exists():
        logging.error(f"Gigapixel executable not found at {GIGAPIXEL_EXE}")
        return

    valid_images = [
        img_path.resolve()
        for img_path in args.images
        if img_path.is_file()
        and img_path.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp", ".webm"}
    ]
    if not valid_images:
        logging.warning("No valid image files found")
        return

    filtered_images = filter_out_upscaled(valid_images, args.force, args.format)
    skipped_count = len(valid_images) - len(filtered_images)
    if skipped_count:
        logging.info(f"Skipping {skipped_count} already upscaled image(s)")

    if not filtered_images:
        logging.warning("No images left to process")
        return

    output_dir = filtered_images[0].parent

    total_images = len(filtered_images)
    completed_images = 0

    def on_progress(path: Path) -> None:
        nonlocal completed_images
        completed_images += 1
        logging.info(f"Saved file: {path.name} [{completed_images}/{total_images}]")

    logging.info(f"Processing {total_images} image(s)")
    try:
        upscale_images(
            filtered_images,
            output_dir,
            args.format,
            on_progress,
        )
        logging.info("All images upscaled successfully")
    except subprocess.CalledProcessError as e:
        logging.error(f"Error processing images: {e.stderr}")


if __name__ == "__main__":
    main()