# coding=utf-8
"""推理加速：自动选择 CUDA / OpenVINO / CPU，并解析已导出的 OpenVINO 模型目录。"""
from __future__ import annotations

import os
from pathlib import Path

_ROOT = Path(__file__).resolve().parents[1]


def project_root() -> Path:
    return _ROOT


def resolve_yolo_weights(pt_path: str | Path | None = None) -> str:
    """优先使用已导出的 OpenVINO 模型，否则回退到 .pt。

    查找顺序:
      1) 环境变量 DISH_YOLO_MODEL
      2) best_openvino_model/（ultralytics export 默认目录）
      3) 传入的 pt_path / best.pt
    """
    env = (os.environ.get("DISH_YOLO_MODEL") or "").strip()
    if env and Path(env).exists():
        return str(Path(env).resolve())

    ov_dir = _ROOT / "best_openvino_model"
    if ov_dir.is_dir() and any(ov_dir.glob("*.xml")):
        return str(ov_dir.resolve())

    if pt_path:
        p = Path(pt_path)
        if p.exists():
            return str(p.resolve())

    fallback = _ROOT / "best.pt"
    return str(fallback.resolve())


def torch_device():
    import torch

    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


def openvino_available() -> bool:
    try:
        import openvino  # noqa: F401
        return True
    except Exception:
        return False


def resolve_feature_openvino() -> Path | None:
    """ConvNeXt 特征提取 OpenVINO 模型目录（可选）。"""
    env = (os.environ.get("DISH_FEATURE_OPENVINO") or "").strip()
    if env:
        p = Path(env)
        if p.is_dir() and any(p.glob("*.xml")):
            return p.resolve()
    default = _ROOT / "feature_openvino_model"
    if default.is_dir() and any(default.glob("*.xml")):
        return default.resolve()
    return None