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

import argparse
import sys
from pathlib import Path
from typing import Iterable

from PIL import Image


SUPPORTED_EXTS = {".jpg", ".jpeg", ".png"}


def _iter_images(path: Path, recursive: bool) -> Iterable[Path]:
    if path.is_file():
        if path.suffix.lower() in SUPPORTED_EXTS:
            yield path
        return

    if not path.is_dir():
        return

    pattern = "**/*" if recursive else "*"
    for p in path.glob(pattern):
        if p.is_file() and p.suffix.lower() in SUPPORTED_EXTS:
            yield p


def _safe_rgb(im: Image.Image, background_rgb: tuple[int, int, int]) -> Image.Image:
    if im.mode == "RGB":
        return im
    if im.mode in {"RGBA", "LA"} or (im.mode == "P" and "transparency" in im.info):
        rgba = im.convert("RGBA")
        bg = Image.new("RGBA", rgba.size, background_rgb + (255,))
        return Image.alpha_composite(bg, rgba).convert("RGB")
    return im.convert("RGB")


def convert_one(
    src: Path,
    dst: Path,
    *,
    dpi: int,
    overwrite: bool,
    background: tuple[int, int, int],
) -> bool:
    if dst.exists() and not overwrite:
        return False

    dst.parent.mkdir(parents=True, exist_ok=True)

    with Image.open(src) as im:
        rgb = _safe_rgb(im, background)
        rgb.save(
            dst,
            format="TIFF",
            compression="tiff_deflate",
            dpi=(dpi, dpi),
        )
    return True


def main(argv: list[str]) -> int:
    parser = argparse.ArgumentParser(
        description="批量将 JPG/PNG 转为 TIFF，并写入 >300 的 DPI（默认 600）。"
    )
    parser.add_argument("input", help="输入文件或目录（目录可批量转换）")
    parser.add_argument(
        "-o",
        "--output",
        default=None,
        help="输出目录（默认：<input_dir>_tiff；若 input 是文件则默认同目录）",
    )
    parser.add_argument(
        "--dpi",
        type=int,
        default=600,
        help="输出 TIFF 的 DPI（必须 >300；默认 600）",
    )
    parser.add_argument(
        "-r",
        "--recursive",
        action="store_true",
        help="输入为目录时递归遍历子目录",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="允许覆盖已存在的输出 tiff 文件",
    )
    parser.add_argument(
        "--background",
        default="255,255,255",
        help="PNG 带透明通道时的底色 RGB，例如 255,255,255（默认白色）",
    )

    args = parser.parse_args(argv)

    if args.dpi <= 300:
        print("错误：--dpi 必须大于 300。", file=sys.stderr)
        return 2

    try:
        bg_parts = [int(x.strip()) for x in str(args.background).split(",")]
        if len(bg_parts) != 3 or any(x < 0 or x > 255 for x in bg_parts):
            raise ValueError
        background = (bg_parts[0], bg_parts[1], bg_parts[2])
    except Exception:
        print("错误：--background 格式应为 r,g,b 且每个值在 0..255。", file=sys.stderr)
        return 2

    in_path = Path(args.input).expanduser()
    if not in_path.exists():
        print(f"错误：输入路径不存在：{in_path}", file=sys.stderr)
        return 2

    out_root: Path | None
    if args.output is None:
        if in_path.is_file():
            out_root = in_path.parent
        else:
            out_root = in_path.with_name(in_path.name + "_tiff")
    else:
        out_root = Path(args.output).expanduser()

    images = list(_iter_images(in_path, args.recursive))
    if not images:
        print("未找到任何 jpg/jpeg/png 文件。")
        return 0

    converted = 0
    skipped = 0
    for src in images:
        if in_path.is_dir():
            rel = src.relative_to(in_path)
            dst = out_root / rel.with_suffix(".tiff")
        else:
            dst = out_root / src.with_suffix(".tiff").name

        ok = convert_one(
            src,
            dst,
            dpi=args.dpi,
            overwrite=args.overwrite,
            background=background,
        )
        if ok:
            converted += 1
        else:
            skipped += 1

    print(f"完成：转换 {converted} 个，跳过 {skipped} 个。输出：{out_root}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))