# -*- coding: utf-8 -*-
"""目录批量抠图：路径校验 / 目录扫描 / 输出文件命名。

三条硬约束（对应需求）：

1. **校验阶段只读不写**：输入目录、输出目录都必须**已经存在**，本模块
   绝不 ``mkdir``。路径写错时直接报错，而不是顺手把目录树建出来 ——
   这样一次误输入（或网页上的恶意路径）不会在磁盘上留下任何东西。
2. **命名固定**：``RMBG_<原主干>_<Unix 时间戳>.<后缀>``，原主干超过
   :data:`STEM_MAXLEN` 个字符时截断。主干统一走
   :func:`birefnet_web.imageops.safe_stem` 清洗，杜绝 ``..`` 之类的穿越。
3. **后缀自适应**：透明输出需要 alpha 通道，``.jpg`` / ``.bmp`` 装不下，
   此时退化为 ``.png``（强行按原后缀保存会被 PIL 静默压成白底图）。

本模块不依赖 torch，可以脱离推理环境单独测试。
"""

from __future__ import annotations

import os
from pathlib import Path
from typing import Iterable, List, Optional, Sequence

from .fsutil import ALPHA_SUFFIXES, IMAGE_SUFFIXES, safe_stem

#: 输出文件名统一前缀（扫描输入目录时也用来自我过滤，避免把产物再抠一遍）
BATCH_PREFIX = "RMBG_"

#: 原文件名主干的最大长度（超过即截断）
STEM_MAXLEN = 20

#: 重名避让时最多尝试的序号
_MAX_DEDUP = 999


class BatchPathError(ValueError):
    """路径不合规：不存在 / 不是目录 / 不可写 / 命中系统目录黑名单。"""


# --------------------------------------------------------------------------- #
# 路径解析与校验
# --------------------------------------------------------------------------- #
def _critical_dirs() -> List[Path]:
    """返回当前系统上「不该往里写文件」的目录（仅用于提示级别的高危拦截）。"""
    roots: List[Path] = []
    names = ("SystemRoot", "windir", "ProgramFiles", "ProgramFiles(x86)", "ProgramData")
    for name in names:
        raw = os.environ.get(name)
        if not raw:
            continue
        p = Path(raw)
        if p.is_absolute():
            roots.append(p)
    if os.name != "nt":
        roots.extend(Path(p) for p in ("/bin", "/sbin", "/etc", "/usr", "/boot", "/dev", "/proc", "/sys"))
    return roots


def resolve_user_path(raw: object, base: Path) -> Path:
    """把用户在界面上敲的一行路径规范化成绝对路径。

    Args:
        raw: 原始输入（可能带首尾空格、成对引号、``~`` 或 ``%VAR%``）。
        base: 相对路径的基准目录（本项目的 ``PROJECT_ROOT``）。

    Returns:
        规范化后的绝对路径（**不校验存在性**，也不做任何创建）。

    Raises:
        BatchPathError: 输入为空。
    """
    text = str(raw or "").strip().strip('"').strip("'").strip()
    if not text:
        raise BatchPathError("路径不能为空")
    text = os.path.expandvars(os.path.expanduser(text))
    path = Path(text)
    if not path.is_absolute():
        path = base / path
    return Path(os.path.normpath(str(path)))


def check_input_dir(raw: object, base: Path) -> Path:
    """校验「图片所在目录」，必须已存在且是目录。

    Raises:
        BatchPathError: 路径为空 / 不存在 / 不是目录。
    """
    path = resolve_user_path(raw, base)
    if not path.exists():
        raise BatchPathError(f"输入目录不存在：{path}（本工具不会自动创建目录，请先确认路径）")
    if not path.is_dir():
        raise BatchPathError(f"输入路径不是目录：{path}")
    return path


def check_output_dir(raw: object, base: Path, *, allow_system_dir: bool = False) -> Path:
    """校验「结果输出目录」，必须已存在、是目录且可写。

    Args:
        raw: 用户输入或默认输出目录。
        base: 相对路径基准。
        allow_system_dir: 为 True 时跳过系统目录黑名单（测试用）。

    Raises:
        BatchPathError: 路径为空 / 不存在 / 不是目录 / 命中黑名单 / 不可写。
    """
    path = resolve_user_path(raw, base)
    if not path.exists():
        raise BatchPathError(f"输出目录不存在：{path}（本工具不会自动创建目录，请先手动建好）")
    if not path.is_dir():
        raise BatchPathError(f"输出路径不是目录：{path}")

    if not allow_system_dir:
        if path == Path(path.anchor):  # C:\ 或 /
            raise BatchPathError(f"拒绝把结果写进磁盘根目录：{path}")
        for critical in _critical_dirs():
            if path == critical or critical in path.parents:
                raise BatchPathError(f"拒绝把结果写进系统目录：{path}")

    if not os.access(path, os.W_OK):
        raise BatchPathError(f"输出目录不可写（权限不足）：{path}")
    return path


# --------------------------------------------------------------------------- #
# 目录扫描
# --------------------------------------------------------------------------- #
def scan_images(
    root: Path,
    *,
    recursive: bool = False,
    skip_dirs: Sequence[Path] = (),
) -> List[Path]:
    """扫描目录下的图片文件。

    Args:
        root: 已校验过的图片目录。
        recursive: 是否包含子目录。
        skip_dirs: 需要跳过的目录（通常是输出目录，避免读到自己的产物）。

    Returns:
        排序后的图片路径列表（先按目录、再按文件名，保证同一批次的处理顺序稳定）。
    """
    skipped = [Path(d).resolve() for d in skip_dirs]
    walker: Iterable[Path] = root.rglob("*") if recursive else root.glob("*")

    found: List[Path] = []
    for path in walker:
        try:
            if not path.is_file():
                continue
        except OSError:  # 权限/坏链接
            continue
        if path.suffix.lower() not in IMAGE_SUFFIXES:
            continue
        # 本工具自己的产物（RMBG_ 前缀）不参与批量，避免反复抠同一张图
        if path.name.startswith(BATCH_PREFIX):
            continue
        if skipped:
            parent = path.parent.resolve()
            if any(parent == s or s in parent.parents or parent in s.parents for s in skipped):
                continue
        found.append(path)

    return sorted(found, key=lambda p: (str(p.parent).lower(), p.name.lower()))


# --------------------------------------------------------------------------- #
# 输出命名
# --------------------------------------------------------------------------- #
def choose_suffix(orig_suffix: str, background: str) -> str:
    """决定输出文件后缀：透明模式下如果原后缀装不下 alpha，就退化为 ``.png``。

    Args:
        orig_suffix: 原文件后缀（含点，如 ``.jpg``）。
        background: 背景模式，``transparent`` 或其它（纯色）。

    Returns:
        合法的输出后缀字符串。
    """
    suffix = (orig_suffix or "").lower()
    if background == "transparent":
        return suffix if suffix in ALPHA_SUFFIXES else ".png"
    return suffix if suffix in IMAGE_SUFFIXES else ".png"


def build_output_name(source: Path, timestamp: int, background: str = "transparent") -> str:
    """按 ``RMBG_<主干(≤20)>_<Unix 时间戳>.<后缀>`` 组装输出文件名。

    Args:
        source: 原始图片路径。
        timestamp: Unix 时间戳（秒），由调用方在**处理时**取，便于追溯。
        background: 背景模式，决定后缀是否允许退化为 ``.png``。

    Returns:
        输出文件名（不含目录）。
    """
    stem = safe_stem(source.name, fallback="image", maxlen=STEM_MAXLEN)
    return f"{BATCH_PREFIX}{stem}_{int(timestamp)}{choose_suffix(source.suffix, background)}"


def unique_path(directory: Path, name: str) -> Path:
    """在同一秒内出现同名输出时，用 ``_1`` / ``_2`` 递增避让，绝不覆盖已有文件。

    Args:
        directory: 输出目录（已存在）。
        name: 目标文件名。

    Returns:
        可安全写入的完整路径。

    Raises:
        BatchPathError: 避让序号耗尽（同名文件超过 :data:`_MAX_DEDUP` 个）。
    """
    target = directory / name
    if not target.exists():
        return target
    for i in range(1, _MAX_DEDUP + 1):
        candidate = directory / f"{target.stem}_{i}{target.suffix}"
        if not candidate.exists():
            return candidate
    raise BatchPathError(f"输出目录中同名文件过多，请检查：{name}")