"""B04 지표면 모델 공통 컨텍스트 및 메시 유틸리티.

지면 마스크가 적용된 포인트에서 footprint(외곽), 격자, 프리뷰 격자를 만들고,
GLB/PLY 프리뷰 및 npz 모델을 원자적으로 저장하는 공통 기능을 제공한다.
5개 표현(TIN/DTM/NURBS/implicit/meshfree) 빌더가 이 컨텍스트를 공유한다.
"""

import hashlib
import json
import math
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable

import numpy as np
import trimesh
from scipy import ndimage

from common_util.common_util_atomic import atomic_write_bytes, atomic_write_npz

MODEL_VERSION = 1
MODEL_METHODS = ("tin", "dtm", "nurbs", "implicit", "meshfree")
SOURCE_FILTER_LABELS = {"grid_min_z": "Grid Min-Z", "csf": "CSF", "pmf": "PMF"}
ProgressCallback = Callable[[int], None]

# 대용량 포인트 배치 처리 크기
_BATCH_SIZE = 500_000


def config_signature(config: dict[str, Any]) -> str:
    """지오메트리 원본과 무관한 등고선·스무딩 설정을 제외한 캐시 서명."""
    sig_config = {
        k: v
        for k, v in config.items()
        if not k.startswith("contour_") and not k.startswith("smoothing_")
    }
    encoded = json.dumps(sig_config, sort_keys=True, default=list).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()[:16]


def bounds_dict(bounds: np.ndarray) -> dict[str, list[float]]:
    return {
        "x": [float(bounds[0, 0]), float(bounds[0, 1])],
        "y": [float(bounds[1, 0]), float(bounds[1, 1])],
        "z": [float(bounds[2, 0]), float(bounds[2, 1])],
    }


def scene_vertices(vertices: np.ndarray, bounds: np.ndarray) -> np.ndarray:
    """모델 좌표를 뷰어(Y-up) 좌표계로 변환한다."""
    center = bounds.mean(axis=1)
    result = np.empty((len(vertices), 3), dtype=np.float32)
    result[:, 0] = vertices[:, 0] - center[0]
    result[:, 1] = vertices[:, 2] - center[2]
    result[:, 2] = -(vertices[:, 1] - center[1])
    return result


def height_colors(vertices: np.ndarray) -> np.ndarray:
    """표고에 따른 그라디언트 정점 색상(RGBA)을 만든다."""
    if not len(vertices):
        return np.empty((0, 4), dtype=np.uint8)
    z = vertices[:, 2]
    span = max(float(np.max(z) - np.min(z)), 1e-9)
    t = np.clip((z - np.min(z)) / span, 0.0, 1.0)
    colors = np.empty((len(vertices), 4), dtype=np.uint8)
    colors[:, 0] = np.clip(36 + 190 * t, 0, 255).astype(np.uint8)
    colors[:, 1] = np.clip(86 + 95 * np.sin(t * np.pi), 0, 255).astype(np.uint8)
    colors[:, 2] = np.clip(128 - 80 * t, 0, 255).astype(np.uint8)
    colors[:, 3] = 255
    return colors


def write_glb(path: Path, vertices: np.ndarray, faces: np.ndarray, bounds: np.ndarray) -> None:
    mesh = trimesh.Trimesh(
        vertices=scene_vertices(vertices, bounds),
        faces=np.asarray(faces, dtype=np.int64),
        vertex_colors=height_colors(vertices),
        process=False,
    )
    payload = mesh.export(file_type="glb")
    if not isinstance(payload, bytes):
        raise TypeError("GLB exporter did not return bytes")
    atomic_write_bytes(path, payload)


def write_binary_ply(
    path: Path, vertices: np.ndarray, normals: np.ndarray, bounds: np.ndarray
) -> None:
    verts = scene_vertices(vertices, bounds)
    scene_normals = np.empty_like(normals, dtype=np.float32)
    scene_normals[:, 0] = normals[:, 0]
    scene_normals[:, 1] = normals[:, 2]
    scene_normals[:, 2] = -normals[:, 1]
    colors = height_colors(vertices)
    dtype = np.dtype(
        [
            ("x", "<f4"),
            ("y", "<f4"),
            ("z", "<f4"),
            ("nx", "<f4"),
            ("ny", "<f4"),
            ("nz", "<f4"),
            ("red", "u1"),
            ("green", "u1"),
            ("blue", "u1"),
            ("alpha", "u1"),
        ]
    )
    records = np.empty(len(vertices), dtype=dtype)
    records["x"], records["y"], records["z"] = verts.T
    records["nx"], records["ny"], records["nz"] = scene_normals.T
    records["red"], records["green"], records["blue"], records["alpha"] = colors.T
    header = (
        "ply\nformat binary_little_endian 1.0\n"
        f"element vertex {len(vertices)}\n"
        "property float x\nproperty float y\nproperty float z\n"
        "property float nx\nproperty float ny\nproperty float nz\n"
        "property uchar red\nproperty uchar green\nproperty uchar blue\nproperty uchar alpha\n"
        "end_header\n"
    ).encode("ascii")
    atomic_write_bytes(path, header + records.tobytes())


def grid_faces(rows: int, cols: int) -> np.ndarray:
    """정규 격자의 삼각형 면 인덱스를 만든다."""
    if rows < 2 or cols < 2:
        return np.empty((0, 3), dtype=np.uint32)
    base = np.arange((rows - 1) * (cols - 1), dtype=np.uint32)
    row = base // (cols - 1)
    col = base % (cols - 1)
    top_left = row * cols + col
    faces = np.empty((len(base) * 2, 3), dtype=np.uint32)
    faces[0::2] = np.stack([top_left, top_left + cols, top_left + 1], axis=1)
    faces[1::2] = np.stack([top_left + 1, top_left + cols, top_left + cols + 1], axis=1)
    return faces


def clip_and_compact_mesh(
    vertices: np.ndarray, faces: np.ndarray, valid_vertices: np.ndarray
) -> tuple[np.ndarray, np.ndarray]:
    """footprint 내부 정점만 사용하는 면을 남기고 미사용 정점을 제거한다."""
    if not len(faces):
        return np.empty((0, 3), np.float32), np.empty((0, 3), np.uint32)
    kept_faces = faces[np.all(valid_vertices[faces], axis=1)]
    if not len(kept_faces):
        return np.empty((0, 3), np.float32), np.empty((0, 3), np.uint32)
    used = np.unique(kept_faces)
    remap = np.full(len(vertices), -1, dtype=np.int64)
    remap[used] = np.arange(len(used))
    return vertices[used], remap[kept_faces].astype(np.uint32)


def grid_vertices(x_coords: np.ndarray, y_coords: np.ndarray, z_grid: np.ndarray) -> np.ndarray:
    xx, yy = np.meshgrid(x_coords, y_coords)
    return np.column_stack([xx.ravel(), yy.ravel(), z_grid.ravel()]).astype(np.float32)


def artifact_size(*paths: Path) -> int:
    return int(sum(path.stat().st_size for path in paths if path.exists()))


@dataclass
class TerrainContext:
    """지면 마스크가 적용된 포인트 집합에서 파생 격자·footprint를 캐싱한다."""

    xyz: np.ndarray
    mask: np.ndarray
    bounds: np.ndarray
    config: dict[str, Any]
    _indices: np.ndarray | None = None
    _samples: dict[int, np.ndarray] = field(default_factory=dict)
    _grids: dict[float, tuple[np.ndarray, np.ndarray, np.ndarray]] = field(default_factory=dict)
    _footprint: tuple[float, float, float, np.ndarray] | None = None

    @property
    def source_count(self) -> int:
        return int(np.count_nonzero(self.mask))

    def indices(self) -> np.ndarray:
        if self._indices is None:
            self._indices = np.flatnonzero(self.mask)
        return self._indices

    def sample(self, maximum: int) -> np.ndarray:
        maximum = max(3, int(maximum))
        if maximum in self._samples:
            return self._samples[maximum]
        indices = self.indices()
        if len(indices) > maximum:
            positions = np.linspace(0, len(indices) - 1, maximum, dtype=np.int64)
            indices = indices[positions]
        points = np.asarray(self.xyz[indices], dtype=np.float32)
        self._samples[maximum] = points
        return points

    def footprint(self) -> tuple[float, float, float, np.ndarray]:
        if self._footprint is not None:
            return self._footprint
        resolution = max(float(self.config.get("footprint_resolution_meters", 1.0)), 0.1)
        x_min, x_max = self.bounds[0]
        y_min, y_max = self.bounds[1]
        cols = max(2, int(math.ceil((x_max - x_min) / resolution)) + 1)
        rows = max(2, int(math.ceil((y_max - y_min) / resolution)) + 1)
        occupied = np.zeros((rows, cols), dtype=bool)
        indices = self.indices()
        for start in range(0, len(indices), _BATCH_SIZE):
            points = np.asarray(self.xyz[indices[start : start + _BATCH_SIZE]], dtype=np.float32)
            gx = np.clip(((points[:, 0] - x_min) / resolution).astype(np.int32), 0, cols - 1)
            gy = np.clip(((points[:, 1] - y_min) / resolution).astype(np.int32), 0, rows - 1)
            occupied[gy, gx] = True
        if not occupied.any():
            raise ValueError("기준 필터에 footprint를 만들 포인트가 없습니다.")

        close_cells = max(
            0,
            int(math.ceil(float(self.config.get("footprint_gap_close_meters", 1.0)) / resolution)),
        )
        footprint = occupied
        if close_cells:
            padded = np.pad(footprint, close_cells, mode="constant", constant_values=False)
            padded = ndimage.binary_closing(
                padded, structure=np.ones((3, 3), dtype=bool), iterations=close_cells
            )
            footprint = padded[close_cells:-close_cells, close_cells:-close_cells]
        if bool(self.config.get("keep_largest_footprint", True)):
            labels, component_count = ndimage.label(
                footprint, structure=np.ones((3, 3), dtype=bool)
            )
            if component_count:
                sizes = np.bincount(labels.ravel())
                sizes[0] = 0
                footprint = labels == int(np.argmax(sizes))
        footprint = ndimage.binary_fill_holes(footprint)

        inset_cells = max(
            0, int(math.ceil(float(self.config.get("boundary_inset_meters", 1.0)) / resolution))
        )
        if inset_cells:
            footprint = ndimage.binary_erosion(
                footprint,
                structure=np.ones((3, 3), dtype=bool),
                iterations=inset_cells,
                border_value=0,
            )
        if not footprint.any():
            raise ValueError("외곽 안쪽 기준 적용 후 유효한 footprint가 없습니다.")
        self._footprint = (float(x_min), float(y_min), resolution, footprint)
        return self._footprint

    def contains_xy(self, x: np.ndarray, y: np.ndarray) -> np.ndarray:
        x_min, y_min, resolution, footprint = self.footprint()
        gx = np.floor((np.asarray(x) - x_min) / resolution).astype(np.int64)
        gy = np.floor((np.asarray(y) - y_min) / resolution).astype(np.int64)
        valid = (gx >= 0) & (gx < footprint.shape[1]) & (gy >= 0) & (gy < footprint.shape[0])
        result = np.zeros(np.broadcast(x, y).shape, dtype=bool)
        result[valid] = footprint[gy[valid], gx[valid]]
        return result

    def footprint_metadata(self) -> dict[str, Any]:
        _, _, resolution, footprint = self.footprint()
        return {
            "footprint_area_m2": round(float(footprint.sum()) * resolution * resolution, 3),
            "footprint_resolution_meters": resolution,
            "boundary_inset_meters": float(self.config.get("boundary_inset_meters", 1.0)),
        }

    def grid(self, resolution: float) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        resolution = max(float(resolution), 0.05)
        cached = self._grids.get(resolution)
        if cached is not None:
            return cached
        x_min, x_max = self.bounds[0]
        y_min, y_max = self.bounds[1]
        cols = max(2, int(math.ceil((x_max - x_min) / resolution)) + 1)
        rows = max(2, int(math.ceil((y_max - y_min) / resolution)) + 1)
        grid = np.full((rows, cols), np.inf, dtype=np.float32)
        indices = self.indices()
        for start in range(0, len(indices), _BATCH_SIZE):
            points = np.asarray(self.xyz[indices[start : start + _BATCH_SIZE]], dtype=np.float32)
            gx = np.clip(((points[:, 0] - x_min) / resolution).astype(np.int32), 0, cols - 1)
            gy = np.clip(((points[:, 1] - y_min) / resolution).astype(np.int32), 0, rows - 1)
            np.minimum.at(grid, (gy, gx), points[:, 2])
        missing = ~np.isfinite(grid)
        if missing.all():
            raise ValueError("기준 필터에 지면 포인트가 없습니다.")
        if missing.any():
            nearest = ndimage.distance_transform_edt(
                missing, return_distances=False, return_indices=True
            )
            grid = grid[tuple(nearest)]
        x_coords = np.linspace(x_min, x_max, cols, dtype=np.float32)
        y_coords = np.linspace(y_min, y_max, rows, dtype=np.float32)
        result = (x_coords, y_coords, grid)
        self._grids[resolution] = result
        return result

    def preview_grid(
        self, preferred_resolution: float
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        x_span = max(float(self.bounds[0, 1] - self.bounds[0, 0]), preferred_resolution)
        y_span = max(float(self.bounds[1, 1] - self.bounds[1, 0]), preferred_resolution)
        maximum = max(4, int(self.config["max_preview_vertices"]))
        predicted = (x_span / preferred_resolution + 1) * (y_span / preferred_resolution + 1)
        if predicted > maximum:
            preferred_resolution *= math.sqrt(predicted / maximum)
        return self.grid(preferred_resolution)

    def clear_caches(self) -> None:
        self._samples.clear()
        self._grids.clear()
        self._indices = None


def with_footprint(context: TerrainContext, metadata: dict[str, Any]) -> dict[str, Any]:
    metadata.update(context.footprint_metadata())
    return metadata


# 모델 빌더가 사용하는 원자적 저장 래퍼 (공통 유틸 재노출)
atomic_npz = atomic_write_npz