# pyright: reportCallIssue=none
# Mitsuba 3 は variant ベースの動的属性解決を行うため Pylance/Pyright が
# mi.* を object 型と推論し reportCallIssue を誤検知する（coding-style.md の方針）．
"""Mitsuba レンダリング実行モジュール．

variant 設定と mi.render() の実行を担う．
`mitsuba.set_variant` はこのモジュールで 1 度だけ呼ぶ．
"""

from __future__ import annotations

import logging
import shutil
from pathlib import Path

import numpy

from sm_tias_lightsim.runtime.llvm_init import ensure_drjit_libllvm_path, log_variant_status
from sm_tias_lightsim.runtime.overrides import apply_drjit_thread_count
from sm_tias_lightsim.scene.builder import build_scene
from sm_tias_lightsim.scene.schema import SceneConfig

_logger = logging.getLogger(__name__)

# 現在設定されている variant を追跡する（プロセス内 1 度だけ変更可能）
_current_variant: str | None = None


def ensure_variant(variant: str) -> None:
    """Mitsuba の variant をプロセス内で 1 度だけ設定する（2 回目以降は何もしない）．

    llvm_* variant 要求時は mi.set_variant 前に DRJIT_LIBLLVM_PATH を自動探索して設定する．
    mitsuba は関数内でインポートするため，モジュール import 時点では LLVM 初期化が走らない．
    set_variant 成功後に variant 状態を INFO ログで出力する．
    既に別の variant が設定されている場合は警告ログを出して続行する
    （Mitsuba は同一プロセスで variant を切り替えられないため）．

    レンダ前だけでなく，レイキャストによる検証（遮蔽チェック等）を目的の variant で
    走らせたいときにも先に呼べるよう，render_scene から切り出してある．

    Args:
        variant: 設定したい Mitsuba variant（例: "llvm_ad_rgb" / "cuda_ad_rgb"）．

    Raises:
        ImportError: JIT の初期化に失敗した場合（LLVM-C.dll 未検出等）．
    """
    global _current_variant  # noqa: PLW0603

    if _current_variant is None:
        # llvm_* variant の場合のみ DRJIT_LIBLLVM_PATH を自動探索する
        # scalar_* / cuda_* は LLVM-C.dll 不要なので探索しない
        # mitsuba の import 前に os.environ['DRJIT_LIBLLVM_PATH'] を設定することが必須．
        # mitsuba を先に import すると LLVM 初期化が走り，LLVM-C.dll 未設定で ImportError になる．
        if variant.startswith("llvm_"):
            ensure_drjit_libllvm_path()

    # mitsuba は DRJIT_LIBLLVM_PATH 設定後にインポートする．
    # モジュールレベルで import すると，ensure_drjit_libllvm_path() の前に LLVM 初期化が走り
    # ImportError になるため，この位置に遅延 import する．
    import mitsuba as mi

    # 共用サーバで LLVM バックエンドが全コアを占有しないよう，SM_DRJIT_THREADS があれば制限する
    apply_drjit_thread_count()

    if _current_variant is None:
        try:
            mi.set_variant(variant)
        except ImportError as e:
            if variant.startswith("llvm_") and "LLVM" in str(e):
                _logger.error(
                    "Mitsuba variant=%s の JIT 初期化に失敗しました．"
                    "LLVM-C.dll が見つかりません．"
                    "Windows: C:\\Program Files\\LLVM\\bin\\LLVM-C.dll を配置するか，"
                    "DRJIT_LIBLLVM_PATH 環境変数を設定してください．"
                    "詳細は docs/02_ENV/ENV_02_環境構築手順.md を参照．",
                    variant,
                )
            raise

        _current_variant = variant
        log_variant_status(_logger, variant, mi.variant())
    elif _current_variant != variant:
        _logger.warning(
            "Mitsuba variant is already set to %r; "
            "requested variant %r will be ignored. "
            "Restart the process to change the variant.",
            _current_variant,
            variant,
        )
    # 同じ variant が既に設定されている場合は何もしない


def render_scene(config: SceneConfig, project_root: Path | None = None) -> numpy.ndarray:
    """SceneConfig からシーンを構築してレンダリングし，HDR 画像（H, W, 3）を返す．

    Mitsuba の variant を config.render.variant に設定する（プロセス内初回のみ有効）．
    既に異なる variant が設定されている場合は警告ログを出して続行する．

    llvm_* variant 要求時は mi.set_variant 前に DRJIT_LIBLLVM_PATH を自動探索して設定する．
    mitsuba は関数内でインポートするため，モジュール import 時点では LLVM 初期化が走らない．
    set_variant 成功後に variant 状態を INFO ログで出力する．

    Args:
        config: 検証済み SceneConfig．
        project_root: STL 相対パスの基点ディレクトリ．None なら CWD を使用．

    Returns:
        HDR 画像 (H, W, 3) の float32 配列．値域は物理量（0〜∞ も許容）．
    """
    ensure_variant(config.render.variant)

    # ensure_variant で DRJIT_LIBLLVM_PATH 設定済み・variant 設定済みのため，ここで import する
    import mitsuba as mi

    scene_dict, tempdir = build_scene(config, project_root=project_root)
    try:
        scene = mi.load_dict(scene_dict)
        image_tensor = mi.render(scene, spp=config.render.spp, seed=config.render.seed)
        arr = numpy.array(image_tensor)
    finally:
        shutil.rmtree(tempdir, ignore_errors=True)

    return arr