"""環境変数による実行時設定の上書き（variant・Dr.Jit スレッド数）．

同じシーン YAML を手元（Windows・llvm_ad_rgb）と研究室サーバ（cuda_ad_rgb）の両方で
使い回すため，YAML を書き換えずに実行環境側から設定を差し替える手段を提供する．
スイープスクリプトは sm-render をサブプロセスで呼ぶので，環境変数にしておけば
`sm-remote run` で 1 回指定するだけで配下の全レンダに伝わる．

- `SM_RENDER_VARIANT`: sm-render が使う Mitsuba variant（YAML の render.variant より優先）
- `SM_DRJIT_THREADS`: Dr.Jit（LLVM バックエンド / scalar レンダ）の CPU スレッド数上限．
  共用サーバ（64 スレッド）で llvm_ad_rgb が全コアを占有するのを防ぐ
  （2026-09-13 Newton で loadavg 33 → 87 を実測）

drjit は `apply_drjit_thread_count` の中で必要な時だけ import する
（未設定時は import しない＝テスト・--help の軽量性を保つ）．
"""

from __future__ import annotations

import logging
import os
from collections.abc import Mapping

_logger = logging.getLogger(__name__)

VARIANT_ENV_VAR = "SM_RENDER_VARIANT"
THREADS_ENV_VAR = "SM_DRJIT_THREADS"

# 同一プロセスで何度呼ばれても INFO ログは 1 回にする（runner と baker の両方から呼ばれるため）
_applied_thread_count: int | None = None


def resolve_render_variant(
    yaml_variant: str,
    cli_variant: str | None = None,
    environ: Mapping[str, str] | None = None,
) -> str:
    """実際に使う Mitsuba variant を決める．

    優先順位: CLI 引数 > 環境変数 `SM_RENDER_VARIANT` > YAML の render.variant．
    空文字・空白のみの値は「指定なし」とみなす．

    Args:
        yaml_variant: シーン YAML の render.variant．
        cli_variant: CLI の --variant（未指定なら None）．
        environ: 参照する環境変数（None なら os.environ）．

    Returns:
        使用する variant 文字列．
    """
    if cli_variant is not None and cli_variant.strip():
        return cli_variant.strip()
    env = os.environ if environ is None else environ
    env_variant = env.get(VARIANT_ENV_VAR, "").strip()
    if env_variant:
        return env_variant
    return yaml_variant


def resolve_thread_count(environ: Mapping[str, str] | None = None) -> int | None:
    """環境変数 `SM_DRJIT_THREADS` から Dr.Jit のスレッド数上限を読む．

    Args:
        environ: 参照する環境変数（None なら os.environ）．

    Returns:
        スレッド数（1 以上）．未設定・空なら None（Dr.Jit 既定＝全論理コア）．

    Raises:
        ValueError: 値が整数でない，または 1 未満の場合．
    """
    env = os.environ if environ is None else environ
    raw = env.get(THREADS_ENV_VAR, "").strip()
    if not raw:
        return None
    try:
        count = int(raw)
    except ValueError as exc:
        msg = f"環境変数 {THREADS_ENV_VAR}={raw!r} は整数ではありません（1 以上の整数を指定）．"
        raise ValueError(msg) from exc
    if count < 1:
        msg = f"環境変数 {THREADS_ENV_VAR}={count} は 1 以上である必要があります．"
        raise ValueError(msg)
    return count


def apply_drjit_thread_count(environ: Mapping[str, str] | None = None) -> int | None:
    """`SM_DRJIT_THREADS` が設定されていれば Dr.Jit のスレッドプールをその数に制限する．

    未設定なら何もしない（drjit も import しない）．`dr.set_thread_count` は
    JIT 初期化の前後どちらで呼んでもよく，冪等なので複数箇所から呼んでよい．

    Args:
        environ: 参照する環境変数（None なら os.environ）．

    Returns:
        設定したスレッド数．未設定なら None．

    Raises:
        ValueError: 環境変数の値が不正な場合．
    """
    global _applied_thread_count  # noqa: PLW0603

    count = resolve_thread_count(environ)
    if count is None:
        return None

    import drjit as dr

    dr.set_thread_count(count)
    if _applied_thread_count != count:
        _logger.info("Dr.Jit の CPU スレッド数を %d に制限しました（%s）", count, THREADS_ENV_VAR)
        _applied_thread_count = count
    return count