"""LLVM JIT 自動検出と診断ログモジュール．

Mitsuba 3 の llvm_* variant を使う際に必要な DRJIT_LIBLLVM_PATH 環境変数を
自動設定する機能と，variant 設定後の状態を診断ログに出力する機能を提供する．

mitsuba は import しない（テストの軽量化・副作用なし設計）．
os.environ への書き込みのみが副作用．
"""

from __future__ import annotations

import glob
import logging
import os
import re
from pathlib import Path

_logger = logging.getLogger(__name__)

# Debian / Ubuntu / Fedora 系は LLVM-C API を単一の libLLVM-<ver>.so（または libLLVM.so.<ver>）に
# 同梱して配布し，libLLVM-C.so を置かない．Dr.Jit 自身はこれらを dlopen できるため，
# 探索対象に含めないと「正常動作しているのに見つからない」誤警告になる（2026-09-13 Newton で確認）．
_POSIX_VERSIONED_LIB_PATTERNS: tuple[str, ...] = (
    "/usr/lib/llvm-*/lib/libLLVM.so",
    "/usr/lib/llvm-*/lib/libLLVM-*.so",
    "/usr/lib/llvm-*/lib/libLLVM-*.so.*",
    "/usr/lib/llvm-*/lib/libLLVM.so.*",
    "/usr/lib/*-linux-gnu/libLLVM-*.so",
    "/usr/lib/*-linux-gnu/libLLVM-*.so.*",
    "/usr/lib/*-linux-gnu/libLLVM.so.*",
    "/usr/lib64/libLLVM-*.so",
    "/usr/lib64/libLLVM.so",
    "/usr/lib64/libLLVM.so.*",
)

_LLVM_VERSION_RE = re.compile(r"(?:llvm-|libLLVM-|libLLVM\.so\.)(\d+)")


def _llvm_version_sort_key(path: str) -> tuple[int, int, str]:
    """LLVM 共有ライブラリのパスを「新しいバージョン優先・.so 優先」で並べるキーを返す．

    パス中の `llvm-<ver>` / `libLLVM-<ver>` / `libLLVM.so.<ver>` のうち最大の数値を
    バージョンとみなす．バージョンが読めないパスは最後に回す．

    Args:
        path: ライブラリのパス文字列．

    Returns:
        sorted() 用のキー（昇順ソートで新しいバージョンが先頭に来る）．
    """
    versions = [int(v) for v in _LLVM_VERSION_RE.findall(path)]
    version = max(versions) if versions else -1
    # 同一バージョン内では開発用シンボリックリンク（*.so）を実体（*.so.N）より優先する
    is_suffixed = 0 if path.endswith(".so") else 1
    return (-version, is_suffixed, path)


def _glob_posix_versioned_libs() -> list[str]:
    """Debian / Ubuntu / Fedora 系のバージョン付き libLLVM を探索して返す．

    Returns:
        見つかったパスのリスト（新しいバージョン優先・重複なし）．
    """
    found: set[str] = set()
    for pattern in _POSIX_VERSIONED_LIB_PATTERNS:
        found.update(glob.glob(pattern))
    return sorted(found, key=_llvm_version_sort_key)


def _iter_default_candidates() -> list[str]:
    """LLVM-C.dll / libLLVM-C.so / libLLVM-*.so の標準インストール先パス候補を返す．

    OS に応じて Windows / POSIX を分岐する．
    テスト容易化のため独立関数として定義する．

    Returns:
        パス候補文字列のリスト（優先度順・重複なし）．
    """
    candidates: list[str] = []

    # 環境変数 LLVM_DIR が設定されていれば bin/ 配下を優先候補に追加する
    llvm_dir = os.environ.get("LLVM_DIR")
    if llvm_dir:
        if os.name == "nt":
            candidates.append(str(Path(llvm_dir) / "bin" / "LLVM-C.dll"))
        else:
            candidates.append(str(Path(llvm_dir) / "lib" / "libLLVM-C.so"))
            candidates.append(str(Path(llvm_dir) / "lib" / "libLLVM-C.dylib"))

    if os.name == "nt":
        # Windows 標準インストール先
        candidates += [
            r"C:\Program Files\LLVM\bin\LLVM-C.dll",
            r"C:\Program Files (x86)\LLVM\bin\LLVM-C.dll",
            r"C:\ProgramData\chocolatey\lib\llvm\tools\LLVM\bin\LLVM-C.dll",
            r"C:\msys64\mingw64\bin\LLVM-C.dll",
            r"C:\msys64\clang64\bin\LLVM-C.dll",
        ]
    else:
        # POSIX 系（Linux / macOS）
        # glob で llvm-* の複数バージョンを一括検索
        candidates += sorted(
            glob.glob("/usr/lib/llvm-*/lib/libLLVM-C.so"), key=_llvm_version_sort_key
        )
        # Debian / Ubuntu（例: /usr/lib/x86_64-linux-gnu/libLLVM-15.so）/ Fedora 系
        candidates += _glob_posix_versioned_libs()
        candidates += [
            "/usr/local/lib/libLLVM-C.so",
            "/opt/homebrew/opt/llvm/lib/libLLVM-C.dylib",
            "/usr/local/opt/llvm/lib/libLLVM-C.dylib",
        ]

    # `libLLVM-*.so` のパターンは libLLVM-C.so にも一致して同じパスが 2 回入るため，
    # 最初の出現（＝優先度の高い位置）だけを残す
    return list(dict.fromkeys(candidates))


def ensure_drjit_libllvm_path(extra_candidates: list[str] | None = None) -> str | None:
    """DRJIT_LIBLLVM_PATH が未設定なら標準パスを探索して os.environ に設定する．

    優先順位：
      1. 既に os.environ['DRJIT_LIBLLVM_PATH'] が設定されている場合はそれを尊重して
         何もしない（戻り値はその既存値）．
      2. 環境変数 LLVM_DIR が設定されていれば `$LLVM_DIR/bin/LLVM-C.dll`（Windows）
         または `$LLVM_DIR/lib/libLLVM-C.so` 等（POSIX）を試す．
      3. Windows の標準インストール先：
         - C:\\Program Files\\LLVM\\bin\\LLVM-C.dll
         - C:\\Program Files (x86)\\LLVM\\bin\\LLVM-C.dll
         - C:\\ProgramData\\chocolatey\\lib\\llvm\\tools\\LLVM\\bin\\LLVM-C.dll
         - C:\\msys64\\mingw64\\bin\\LLVM-C.dll
         - C:\\msys64\\clang64\\bin\\LLVM-C.dll
      4. POSIX 系：
         - /usr/lib/llvm-*/lib/libLLVM-C.so（glob）
         - Debian / Ubuntu / Fedora 系のバージョン付き libLLVM（glob．新しいバージョン優先）:
           /usr/lib/llvm-*/lib/libLLVM{,-*}.so*，/usr/lib/*-linux-gnu/libLLVM-*.so*，
           /usr/lib/*-linux-gnu/libLLVM.so.*，/usr/lib64/libLLVM*.so*
         - /usr/local/lib/libLLVM-C.so
         - /opt/homebrew/opt/llvm/lib/libLLVM-C.dylib
         - /usr/local/opt/llvm/lib/libLLVM-C.dylib
      5. extra_candidates が指定されていればそれを最後に試す．

    Args:
        extra_candidates: 追加で試すパス候補のリスト（任意）．
            標準候補で見つからない場合に最後に試す．

    Returns:
        発見・設定したパス文字列．既設定だった場合はその値．
        何も見つからなければ None．

    Side effect:
        os.environ['DRJIT_LIBLLVM_PATH'] を設定する（未設定かつ発見時のみ）．
    """
    # 1. 既に設定されている場合はそのまま返す
    existing = os.environ.get("DRJIT_LIBLLVM_PATH")
    if existing:
        _logger.debug("DRJIT_LIBLLVM_PATH は既に設定されています: %s", existing)
        return existing

    # 標準候補 + 追加候補を順番に探索する
    all_candidates = _iter_default_candidates()
    if extra_candidates:
        all_candidates += extra_candidates

    for candidate in all_candidates:
        if Path(candidate).is_file():
            os.environ["DRJIT_LIBLLVM_PATH"] = candidate
            _logger.info(
                "DRJIT_LIBLLVM_PATH を自動設定しました: %s",
                candidate,
            )
            return candidate

    _logger.warning(
        "LLVM-C.dll / libLLVM-C.so / libLLVM-*.so が標準パスに見つかりませんでした．"
        "LLVM をインストールするか，DRJIT_LIBLLVM_PATH 環境変数を手動で設定してください．"
        "詳細は docs/02_ENV/ENV_02_環境構築手順.md を参照してください．"
    )
    return None


def log_variant_status(logger: logging.Logger, requested_variant: str, actual_variant: str) -> None:
    """variant 設定後の状態を INFO / WARNING でログ出力する．

    actual_variant が requested_variant と異なる（scalar fallback 等）場合は
    WARNING で明示する．

    Args:
        logger: 出力先の Logger オブジェクト．
        requested_variant: 要求した variant 文字列（YAML の variant フィールド等）．
        actual_variant: `mi.variant()` で確認した実際の variant 文字列．
    """
    if actual_variant.startswith("llvm_"):
        llvm_path = os.environ.get("DRJIT_LIBLLVM_PATH", "（未設定）")
        logger.info(
            "Mitsuba variant=%s, LLVM JIT 経路で実行（DRJIT_LIBLLVM_PATH=%s）",
            actual_variant,
            llvm_path,
        )
    elif actual_variant.startswith("cuda_"):
        logger.info("Mitsuba variant=%s, CUDA GPU 経路で実行", actual_variant)
    elif actual_variant.startswith("scalar_"):
        logger.info("Mitsuba variant=%s, scalar 経路で実行（単スレッド）", actual_variant)
    else:
        logger.info("Mitsuba variant=%s", actual_variant)

    # 要求が llvm_* なのに scalar_* に落ちている場合は WARNING を追加する
    if requested_variant.startswith("llvm_") and actual_variant.startswith("scalar_"):
        logger.warning(
            "要求 variant=%s に対し実 variant=%s"
            "（LLVM 経路に乗れていません．"
            "LLVM-C.dll の場所と DRJIT_LIBLLVM_PATH を確認してください）",
            requested_variant,
            actual_variant,
        )