#!/usr/bin/env python3
from __future__ import annotations

import argparse
import os
import shutil
import subprocess
import sys
from pathlib import Path


class LocalUpscaler:
    supported_formats = frozenset({".jpg", ".jpeg", ".png", ".webp"})

    def __init__(self, binary: Path, models: Path) -> None:
        self.binary = binary.expanduser().resolve()
        self.models = models.expanduser().resolve()

    @staticmethod
    def default_output(input_path: Path, scale: int) -> Path:
        return input_path.with_name(f"{input_path.stem}-x{scale}.png")

    def upscale(
        self,
        input_path: Path,
        output_path: Path,
        scale: int,
        model: str,
    ) -> None:
        input_path = input_path.expanduser().resolve()
        output_path = output_path.expanduser().resolve()

        if not input_path.is_file():
            raise ValueError(f"input image does not exist: {input_path}")
        if input_path.suffix.lower() not in self.supported_formats:
            raise ValueError("input must be a JPG, PNG, or WebP image")
        if output_path.suffix.lower() not in self.supported_formats:
            raise ValueError("output must use .jpg, .jpeg, .png, or .webp")
        if input_path == output_path:
            raise ValueError("output must be different from the input image")
        if not self.binary.is_file():
            raise ValueError(f"Real-ESRGAN binary not found: {self.binary}")
        if not self.models.is_dir():
            raise ValueError(f"model directory not found: {self.models}")
        for suffix in (".param", ".bin"):
            model_file = self.models / f"{model}{suffix}"
            if not model_file.is_file():
                raise ValueError(f"model file not found: {model_file}")

        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_format = output_path.suffix.lower().lstrip(".")
        if output_format == "jpeg":
            output_format = "jpg"
        subprocess.run(
            [
                str(self.binary),
                "-i",
                str(input_path),
                "-o",
                str(output_path),
                "-n",
                model,
                "-s",
                str(scale),
                "-m",
                str(self.models),
                "-f",
                output_format,
            ],
            check=True,
        )


def default_binary() -> Path:
    configured = os.environ.get("REALESRGAN_BIN")
    if configured:
        return Path(configured)
    executable = "realesrgan-ncnn-vulkan.exe" if os.name == "nt" else "realesrgan-ncnn-vulkan"
    installed = shutil.which(executable)
    return Path(installed) if installed else Path.cwd() / executable


def argument_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Upscale one image locally with Real-ESRGAN NCNN Vulkan."
    )
    parser.add_argument("input", type=Path, help="JPG, PNG, or WebP input image")
    parser.add_argument("--scale", type=int, choices=(2, 4), default=2)
    parser.add_argument("--output", type=Path, help="Output path (default: NAME-xSCALE.png)")
    parser.add_argument("--binary", type=Path, help="Path to realesrgan-ncnn-vulkan")
    parser.add_argument("--models", type=Path, help="Path to the NCNN models directory")
    parser.add_argument("--model", default="realesrgan-x4plus", help="NCNN model name")
    return parser


def main() -> int:
    arguments = argument_parser().parse_args()
    binary = arguments.binary or default_binary()
    configured_models = os.environ.get("REALESRGAN_MODELS")
    models = arguments.models or (
        Path(configured_models) if configured_models else binary.parent / "models"
    )
    output = arguments.output or LocalUpscaler.default_output(
        arguments.input, arguments.scale
    )

    try:
        LocalUpscaler(binary, models).upscale(
            arguments.input,
            output,
            arguments.scale,
            arguments.model,
        )
    except (OSError, ValueError, subprocess.CalledProcessError) as error:
        print(f"upscale_image.py: {error}", file=sys.stderr)
        return 1

    print(output.expanduser().resolve())
    return 0


if __name__ == "__main__":
    raise SystemExit(main())