from __future__ import annotations

from pathlib import Path

import matplotlib
import numpy as np
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

matplotlib.use("Agg")
import matplotlib.pyplot as plt


def _axis_limits_from_points(points: np.ndarray) -> tuple[tuple[float, float], tuple[float, float], tuple[float, float]]:
    x_min, x_max = float(np.min(points[:, 0])), float(np.max(points[:, 0]))
    y_min, y_max = float(np.min(points[:, 1])), float(np.max(points[:, 1]))
    z_min, z_max = float(np.min(points[:, 2])), float(np.max(points[:, 2]))

    max_range = max(x_max - x_min, y_max - y_min, z_max - z_min) / 2.0
    max_range = max(max_range, 1e-6) * 1.08

    x_mid = (x_max + x_min) / 2.0
    y_mid = (y_max + y_min) / 2.0
    z_mid = (z_max + z_min) / 2.0

    return (
        (x_mid - max_range, x_mid + max_range),
        (y_mid - max_range, y_mid + max_range),
        (z_mid - max_range, z_mid + max_range),
    )


def _to_vertices(face_v_item: object) -> np.ndarray | None:
    verts = np.asarray(face_v_item, dtype=float)
    if verts.ndim == 2 and verts.shape[1] == 3 and len(verts) >= 3:
        return verts

    flat = verts.reshape(-1)
    if flat.size >= 9 and flat.size % 3 == 0:
        shaped = flat.reshape(-1, 3)
        if len(shaped) >= 3:
            return shaped

    return None


def visualize_geometry(
    geometry_npz: str | Path,
    output_png: str | Path,
    *,
    elev: float = 45.0,
    azim: float = 15.0,
    dpi: int = 300,
) -> Path:
    """Render geometry polygons from PACK geometry npz (expects key: face_v)."""
    geometry_npz = Path(geometry_npz)
    output_png = Path(output_png)

    with np.load(geometry_npz, allow_pickle=True) as data:
        if "face_v" not in data:
            keys = ", ".join(sorted(data.files))
            raise KeyError(f"Missing key 'face_v' in {geometry_npz}; keys=[{keys}]")
        face_v = data["face_v"]

    polygons: list[np.ndarray] = []
    for item in face_v:
        verts = _to_vertices(item)
        if verts is not None:
            polygons.append(verts)

    if not polygons:
        raise ValueError(f"No valid polygons in {geometry_npz}")

    fig = plt.figure(figsize=(4, 4))
    ax = fig.add_subplot(111, projection="3d")
    ax.view_init(elev=elev, azim=azim)
    fig.subplots_adjust(left=0, right=1, top=1, bottom=0)

    for verts in polygons:
        collection = Poly3DCollection(
            [verts],
            facecolors="#FFFFFF",
            edgecolors="#4D4D4D",
            linewidths=0.42,
            alpha=0.35,
        )
        ax.add_collection3d(collection)

    all_points = np.vstack(polygons)
    x_lim, y_lim, z_lim = _axis_limits_from_points(all_points)
    ax.set_xlim(*x_lim)
    ax.set_ylim(*y_lim)
    ax.set_zlim(*z_lim)
    ax.set_box_aspect([1, 1, 1])
    ax.set_axis_off()

    output_png.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output_png, dpi=dpi, bbox_inches="tight", pad_inches=0.04, transparent=True)
    plt.close(fig)
    return output_png