"""PyVista 3D インタラクティブビューアモジュール．

シーンのメッシュ・LED・センサを PyVista で可視化する．
照度ベイク結果 (irradiance) があればメッシュを照度で着色する．
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np

from sm_tias_lightsim.geometry.loader import load_stl
from sm_tias_lightsim.geometry.transforms import scale_vertices
from sm_tias_lightsim.scene.schema import (
    AreaLight,
    DirectionalLight,
    PointLight,
    SceneConfig,
)

if TYPE_CHECKING:
    import pyvista as pv

_logger = logging.getLogger(__name__)

# YAML 座標 (mm) → m 変換係数
_MM_TO_M = 0.001

# 方向ベクトルの長さがこれ以下なら縮退とみなしてスキップする
_DIR_NORM_EPSILON = 1e-9

# 光源・センサを示す球の半径（m 単位）
_MARKER_SPHERE_RADIUS = 0.005

# AreaLight の放射方向矢印のスケール（m 単位）
_AREA_LIGHT_ARROW_SCALE = 0.02

# DirectionalLight の矢印をシーン対角線長に対して拡大する比率
_DIRECTIONAL_ARROW_SCALE_RATIO = 0.3

# センサの視線矢印をシーン対角線長に対してクランプする比率
_SENSOR_ARROW_SCALE_RATIO = 0.2


def _build_pyvista_polydata(vertices: np.ndarray, faces: np.ndarray) -> pv.PolyData:
    """頂点・面配列から PyVista PolyData を構築する．

    PyVista の faces 形式は各行が [n_vertices, v0, v1, v2, ...] の 1D フラット配列．

    Args:
        vertices: 頂点配列 (N_v, 3)，m 単位．
        faces: 三角形面配列 (N_f, 3)，頂点インデックス．

    Returns:
        PyVista PolyData オブジェクト．
    """
    import pyvista as pv

    n_faces = len(faces)
    # PyVista 形式: [3, v0, v1, v2, 3, v0, v1, v2, ...]
    pv_faces = np.hstack([np.full((n_faces, 1), 3, dtype=np.int64), faces]).flatten()
    return pv.PolyData(vertices.astype(np.float32), pv_faces)


def _get_scene_bbox_center_and_diag(
    scene_config: SceneConfig, project_root: Path
) -> tuple[np.ndarray, float]:
    """シーン全メッシュのバウンディングボックス中心と対角線長を返す．

    DirectionalLight の矢印位置を決めるために使う．
    メッシュが空なら原点と 1.0 を返す．

    Args:
        scene_config: SceneConfig．
        project_root: STL 相対パスの基点ディレクトリ．

    Returns:
        (center (3,), diagonal_length) のタプル．
    """
    all_vertices: list[np.ndarray] = []
    for mesh_cfg in scene_config.geometry.meshes:
        try:
            raw_path = Path(mesh_cfg.path)
            stl_path = raw_path if raw_path.is_absolute() else project_root / mesh_cfg.path
            loaded = load_stl(stl_path, mesh_cfg.unit)
            scaled = scale_vertices(loaded, target_unit="m")
            all_vertices.append(scaled.vertices)
        except Exception as exc:
            _logger.warning("バウンディングボックス計算でメッシュを読み込めませんでした: %s", exc)

    if not all_vertices:
        return np.zeros(3), 1.0

    combined = np.vstack(all_vertices)
    mins = combined.min(axis=0)
    maxs = combined.max(axis=0)
    center = (mins + maxs) / 2.0
    diagonal = float(np.linalg.norm(maxs - mins))
    return center, max(diagonal, 1e-6)


def show_scene(
    scene_config: SceneConfig,
    project_root: Path,
    irradiance: dict[str, np.ndarray] | None = None,
    *,
    off_screen: bool = False,
    screenshot: Path | None = None,
) -> None:
    """PyVista でシーンとベイク結果を表示する．

    メッシュ・LED・センサ・座標軸を含む 3D インタラクティブビューアを開く．
    irradiance が指定されていれば，メッシュを照度で着色する（viridis カラーマップ）．

    off_screen=True のときは GUI ウィンドウを開かず，screenshot が指定されていれば保存する．
    screenshot=None かつ off_screen=True のときは何も出力しない（no-op）．

    Args:
        scene_config: 検証済み SceneConfig．
        project_root: STL 相対パスの基点ディレクトリ．
        irradiance: {mesh_id: irradiance_per_face (N_f,)} の辞書．None なら着色なし．
        off_screen: True のとき GUI ウィンドウを開かない（screenshot 専用）．
        screenshot: 保存先 PNG パス．指定時は off_screen で screenshot を保存する．
    """
    import pyvista as pv

    # screenshot 指定があれば off_screen を強制する
    if screenshot is not None:
        off_screen = True

    plotter = pv.Plotter(off_screen=off_screen)

    # ---- メッシュ追加 ----
    for mesh_cfg in scene_config.geometry.meshes:
        try:
            raw_path = Path(mesh_cfg.path)
            stl_path = raw_path if raw_path.is_absolute() else project_root / mesh_cfg.path
            loaded = load_stl(stl_path, mesh_cfg.unit)
            scaled = scale_vertices(loaded, target_unit="m")
        except Exception as exc:
            _logger.warning("メッシュ '%s' を読み込めませんでした: %s", mesh_cfg.id, exc)
            continue

        polydata = _build_pyvista_polydata(scaled.vertices, scaled.faces)

        if irradiance is not None and mesh_cfg.id in irradiance:
            irr = irradiance[mesh_cfg.id]
            polydata.cell_data["irradiance"] = irr.astype(np.float32)
            plotter.add_mesh(
                polydata,
                scalars="irradiance",
                cmap="viridis",
                show_scalar_bar=True,
                scalar_bar_args={"title": "Irradiance (W/m²)"},
                label=mesh_cfg.id,
            )
        else:
            plotter.add_mesh(polydata, color="lightgray", opacity=0.5, label=mesh_cfg.id)

    # ---- シーン全体のバウンディングボックス（DirectionalLight 用）----
    bbox_center, bbox_diag = _get_scene_bbox_center_and_diag(scene_config, project_root)

    # ---- LED / 光源表示 ----
    for light in scene_config.lights:
        if isinstance(light, PointLight):
            pos_m = [c * _MM_TO_M for c in light.position]
            sphere = pv.Sphere(radius=_MARKER_SPHERE_RADIUS, center=pos_m)
            plotter.add_mesh(sphere, color="yellow", label=f"LED:{light.id}")

        elif isinstance(light, AreaLight):
            center_m = [c * _MM_TO_M for c in light.center]
            size_m = [light.size[0] * _MM_TO_M, light.size[1] * _MM_TO_M]

            dir_vec = np.array(light.direction, dtype=np.float64)
            dir_norm = np.linalg.norm(dir_vec)
            if dir_norm > _DIR_NORM_EPSILON:
                dir_vec = dir_vec / dir_norm

            plane = pv.Plane(
                center=center_m,
                direction=dir_vec.tolist(),
                i_size=size_m[0],
                j_size=size_m[1],
            )
            plotter.add_mesh(plane, color="orange", opacity=0.7, label=f"AreaLight:{light.id}")

            arrow = pv.Arrow(
                start=center_m,
                direction=dir_vec.tolist(),
                scale=_AREA_LIGHT_ARROW_SCALE,
            )
            plotter.add_mesh(arrow, color="orange", label=f"AreaLight:{light.id}:dir")

        elif isinstance(light, DirectionalLight):
            # バウンディングボックス中心から矢印 1 本で進行方向を示す
            dir_vec = np.array(light.direction, dtype=np.float64)
            dir_norm = np.linalg.norm(dir_vec)
            if dir_norm < _DIR_NORM_EPSILON:
                continue
            dir_vec = dir_vec / dir_norm

            arrow_scale = bbox_diag * _DIRECTIONAL_ARROW_SCALE_RATIO
            arrow = pv.Arrow(
                start=bbox_center.tolist(),
                direction=dir_vec.tolist(),
                scale=arrow_scale,
            )
            plotter.add_mesh(arrow, color="lightyellow", label=f"DirLight:{light.id}")

    # ---- センサ表示 ----
    sensor_pos_m = [c * _MM_TO_M for c in scene_config.sensor.position]
    sensor_target_m = [c * _MM_TO_M for c in scene_config.sensor.target]

    sensor_sphere = pv.Sphere(radius=_MARKER_SPHERE_RADIUS, center=sensor_pos_m)
    plotter.add_mesh(sensor_sphere, color="blue", label="sensor")

    sensor_pos_arr = np.array(sensor_pos_m)
    sensor_target_arr = np.array(sensor_target_m)
    view_dir = sensor_target_arr - sensor_pos_arr
    view_len = np.linalg.norm(view_dir)
    if view_len > _DIR_NORM_EPSILON:
        view_dir_n = view_dir / view_len
        sensor_arrow = pv.Arrow(
            start=sensor_pos_m,
            direction=view_dir_n.tolist(),
            scale=min(float(view_len), bbox_diag * _SENSOR_ARROW_SCALE_RATIO),
        )
        plotter.add_mesh(sensor_arrow, color="blue", label="sensor:view")

    # ---- 座標軸と境界表示 ----
    plotter.add_axes()
    plotter.show_bounds(grid="back", location="outer")

    # ---- 表示 / 保存 ----
    if screenshot is not None:
        plotter.show(screenshot=str(screenshot))
    elif not off_screen:
        plotter.show()
    # off_screen=True かつ screenshot=None の場合は no-op