# -*- coding: utf-8 -*-
"""模型代码定位与 ComfyUI 解耦层。

本项目把原来 ComfyUI 节点 ``comfyui_birefnet_ll`` 的模型代码**内联**到了
``vendor/comfyui_birefnet_ll/``，因此默认情况下完全自包含，不再依赖外部 ComfyUI 安装。

定位优先级（``find_node_dir``）：
    1. 命令行 ``--node-dir``
    2. 环境变量 ``BIREFNET_NODE_DIR``
    3. **项目内 ``vendor/comfyui_birefnet_ll``（默认，开箱即用）**
    4. ``config.json`` 中的 ``node_dir`` 字段（兼容旧配置）
    5. 自动扫描本机 ComfyUI 安装（便于跟随上游节点升级）

模型包在 import 阶段依赖 ComfyUI 的 ``folder_paths``，但只用于「按文件名查权重路径」
（即骨干网络预训练权重）。本项目始终以 ``bb_pretrained=False`` 构建骨干网络，不会真正
读取这类权重，因此这里注入一个最小实现的 ``folder_paths`` 垫片即可脱离 ComfyUI 运行。
"""

from __future__ import annotations

import importlib
import os
import sys
import types
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Tuple

# --------------------------------------------------------------------------- #
# 路径常量
# --------------------------------------------------------------------------- #
PROJECT_ROOT = Path(__file__).resolve().parent.parent

#: 项目内自带的模型代码（默认来源，随项目一起分发）
VENDOR_DIR = PROJECT_ROOT / "vendor" / "comfyui_birefnet_ll"

#: 项目内自带的权重目录（默认来源）
LOCAL_MODEL_DIR = PROJECT_ROOT / "models"

#: 环境变量：手动指定模型代码目录
NODE_DIR_ENV = "BIREFNET_NODE_DIR"

#: 节点仓库名，用于兜底扫描外部 ComfyUI 安装
NODE_REPO_NAME = "comfyui_birefnet_ll"

#: 旧版架构权重文件名（与 birefnetNode.py 保持一致）
OLD_MODEL_FILES: Tuple[str, ...] = ("BiRefNet-DIS_ep580.pth", "BiRefNet-ep480.pth")

#: 判定「是有效模型代码目录」的标志文件
_MARKER = ("birefnet", "models", "birefnet.py")

_BUNDLED_PYTHON = PROJECT_ROOT / "python" / "python.exe"

_INSTALL_HINT = (
    "缺少运行依赖。本项目自带独立 Python 运行时，请使用：\n"
    f'  "{_BUNDLED_PYTHON}" Webui.py\n'
    "或双击 run.bat。\n"
    "若想用自己的解释器，请先安装依赖：pip install -r requirements.txt"
)


# --------------------------------------------------------------------------- #
# 模型代码目录定位
# --------------------------------------------------------------------------- #
def _is_code_dir(path: Path) -> bool:
    """目录内是否含 ``birefnet/models/birefnet.py``。"""
    try:
        return path.joinpath(*_MARKER).is_file()
    except OSError:
        return False


def _scan_external_node_dirs() -> List[Path]:
    """扫描本机常见的 ComfyUI 安装位置（仅在需要跟随上游升级时才会命中）。"""
    found: List[Path] = []
    home = Path.home()
    roots: List[Path] = [home]
    for letter in "CDEFGHIJ":
        roots.append(Path(f"{letter}:/"))
    patterns = (
        "ComfyUI/custom_nodes/" + NODE_REPO_NAME,
        "ComfyUI*/ComfyUI/custom_nodes/" + NODE_REPO_NAME,
        "ComfyUI*/custom_nodes/" + NODE_REPO_NAME,
        "ComfyUI*/ComfyUI/custom_nodes/ComfyUI_BiRefNet_ll",
        "ComfyUI*/custom_nodes/ComfyUI_BiRefNet_ll",
        "*/*/ComfyUI/custom_nodes/" + NODE_REPO_NAME,
    )
    for root in roots:
        if not root.exists():
            continue
        for pat in patterns:
            try:
                found.extend(sorted(root.glob(pat)))
            except (OSError, ValueError):
                continue
    return found


def _dedup(paths: Iterable[Path]) -> List[Path]:
    seen, out = set(), []
    for p in paths:
        try:
            key = str(p.resolve()).lower() if p.exists() else str(p).lower()
        except OSError:
            key = str(p).lower()
        if key in seen:
            continue
        seen.add(key)
        out.append(p)
    return out


def candidate_code_dirs(config_hint: Optional[str] = None) -> List[Path]:
    """按优先级生成候选模型代码目录。"""
    cands: List[Path] = []

    env = os.environ.get(NODE_DIR_ENV)
    if env:
        cands.append(Path(env))

    # 项目自带（默认命中）
    cands.append(VENDOR_DIR)

    # 旧配置中记录的路径（兼容早期版本写死的 ComfyUI 节点路径）
    if config_hint:
        cands.append(Path(config_hint))

    # 外部 ComfyUI 安装（兜底）
    cands.append(PROJECT_ROOT.parent / NODE_REPO_NAME)
    cands.extend(_scan_external_node_dirs())

    return _dedup(cands)


def find_node_dir(explicit: Optional[str] = None, config_hint: Optional[str] = None) -> Path:
    """定位模型代码目录（需含 ``birefnet/models/birefnet.py``）。

    正常情况下会直接命中项目内的 ``vendor/comfyui_birefnet_ll``。
    """
    cands: List[Path] = []
    if explicit:
        cands.append(Path(explicit))
    cands.extend(candidate_code_dirs(config_hint))

    for cand in cands:
        if _is_code_dir(cand):
            try:
                return cand.resolve()
            except OSError:
                return cand

    raise FileNotFoundError(
        "未找到 BiRefNet 模型代码目录（应包含 birefnet/models/birefnet.py）。\n"
        f"默认位置：{VENDOR_DIR}\n"
        "若该目录缺失或被误删，可从 ComfyUI 节点 comfyui_birefnet_ll 重新拷贝，或用以下方式指定：\n"
        "  1) 启动参数 --node-dir <路径>\n"
        f"  2) 环境变量 {NODE_DIR_ENV}=<路径>\n"
        "  3) config.json 中的 node_dir 字段\n"
        "已尝试的候选目录：\n  - " + "\n  - ".join(str(c) for c in cands[:8])
    )


def default_model_dirs(code_dir: Optional[Path] = None) -> List[Path]:
    """返回默认权重搜索目录：项目内 ``models/`` 优先，其次外部目录。"""
    dirs: List[Path] = [LOCAL_MODEL_DIR]

    # 若用户显式指向了外部 ComfyUI 节点，则顺带搜索对应的 models/BiRefNet
    if code_dir is not None:
        p = Path(code_dir)
        try:
            if p.resolve() != VENDOR_DIR.resolve():
                dirs.append(p.parent.parent / "models" / "BiRefNet")
        except OSError:
            pass

    return _dedup(dirs)


# --------------------------------------------------------------------------- #
# folder_paths 垫片
# --------------------------------------------------------------------------- #
class _FolderPathsShim(types.ModuleType):
    """仅实现 birefnet/config.py 用到的接口，一律返回 None（不查权重）。"""

    def __init__(self, model_dirs: Iterable[Path]) -> None:
        super().__init__("folder_paths")
        self.models_dir = str(next(iter(model_dirs), Path("models")))
        self.folder_names_and_paths: Dict[str, Tuple[List[str], set]] = {
            "birefnet": ([self.models_dir], {".pt", ".pth", ".safetensors", ".ckpt"})
        }
        self.supported_pt_extensions = {".pt", ".pth", ".safetensors", ".ckpt"}

    def get_folder_paths(self, folder_name: str) -> List[str]:
        return list(self.folder_names_and_paths.get(folder_name, ([], None))[0])

    def get_full_path(self, folder_name: str, filename: str) -> Optional[str]:
        return None

    def get_filename_list(self, folder_name: str) -> List[str]:
        return []

    def add_model_folder_path(self, folder_name: str, full_folder_path: str) -> None:
        if folder_name not in self.folder_names_and_paths:
            self.folder_names_and_paths[folder_name] = ([], set())
        self.folder_names_and_paths[folder_name][0].append(str(full_folder_path))

    def filter_files_extensions(self, files, extensions):  # pragma: no cover
        return [f for f in files if os.path.splitext(f)[1] in extensions]


def ensure_folder_paths(model_dirs: Iterable[Path]) -> bool:
    """确保 ``folder_paths`` 可用。返回 True 表示用了真实实现（即运行在 ComfyUI 内）。"""
    if "folder_paths" in sys.modules:
        return not isinstance(sys.modules["folder_paths"], _FolderPathsShim)
    try:  # 若真的在 ComfyUI 环境内运行，优先用真实实现
        importlib.import_module("folder_paths")
        return True
    except Exception:
        pass
    sys.modules["folder_paths"] = _FolderPathsShim(model_dirs)
    return False


# --------------------------------------------------------------------------- #
# 模型类导入
# --------------------------------------------------------------------------- #
_cached: Dict[str, object] = {}


def load_model_classes(node_dir: Path, model_dirs: Iterable[Path]) -> Dict[str, object]:
    """导入并返回 ``{"BiRefNet":..., "OldBiRefNet":..., "check_state_dict":..., "shim": bool}``。

    结果会被缓存，重复调用不会重复导入。
    """
    if _cached:
        return _cached

    node_dir = Path(node_dir).resolve()
    shim = ensure_folder_paths(model_dirs)

    p = str(node_dir)
    if p not in sys.path:
        sys.path.insert(0, p)

    try:
        from birefnet.models.birefnet import BiRefNet
        from birefnet.utils import check_state_dict
    except ImportError as exc:  # pragma: no cover - 环境问题
        raise ImportError(f"导入 birefnet 模型包失败（{node_dir}）：{exc}\n{_INSTALL_HINT}") from exc

    try:
        from birefnet_old.models.birefnet import BiRefNet as OldBiRefNet
    except Exception:  # 旧包缺失不影响新模型使用
        OldBiRefNet = None  # type: ignore[assignment]

    _cached.update(
        BiRefNet=BiRefNet,
        OldBiRefNet=OldBiRefNet,
        check_state_dict=check_state_dict,
        node_dir=node_dir,
        used_folder_paths_shim=shim,
    )
    return _cached