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

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

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

import numpy as np
import trimesh
from scipy import ndimage
from skimage.restoration import inpaint_biharmonic

from common_util.common_util_atomic import atomic_write_bytes, atomic_write_npz

logger = logging.getLogger(__name__)

MODEL_VERSION = 1
ProgressCallback = Callable[[int], None]

# 대용량 포인트 배치 처리 크기
_BATCH_SIZE = 500_000
# implicit 격자 칸 뽑기 씨앗 — 고정이라 다시 지어도 같은 중심점
IMPLICIT_SAMPLE_SEED = 558

# 서명에서 빼는 키 — 지오메트리를 만드는 방법이 아니라 "무엇을 만들지"를 고르는 값들.
# source_filters·precompute를 넣으면 조합을 바꿀 때마다 manifest가 통째로 폐기돼
# 이전에 만들어 둔 조합이 사라진다. manifest는 이미 필터·방식별로 키가 갈린다.
_SIGNATURE_EXCLUDED_KEYS = {"source_filters", "precompute"}


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_")
        and k not in _SIGNATURE_EXCLUDED_KEYS
    }
    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 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 fill_holes_smooth(
    z: np.ndarray, known: np.ndarray, inside: np.ndarray, min_gap_cells: float
) -> np.ndarray:
    """외곽 안 큰 구멍을 둘레 지면에서 매끈하게(쌍조화 · 판 굽힘) 이어 메움(766-2).

    구멍 = 땅 칸에서 `min_gap_cells` 넘게 떨어진 칸 + 그 둘레 `min_gap_cells` 안 빈 칸(외곽 안) ·
    작은 틈은 지금처럼 가까운 칸 값 · 둘레 높이 · 기울기를 이어 받음(평면 · 곡면 비탈 그대로) ·
    외곽 밖 이웃도 같이 풀어 가까운 칸 값에 끌리지 않게 함.
    돌려줌 = 땅 칸에서 `min_gap_cells` 넘게 먼 메운 칸(메운 넓이 지표 · implicit 길잡이).
    """
    gap = ~known & inside
    far = (ndimage.distance_transform_edt(~known) > min_gap_cells) & gap
    if not far.any():
        return far
    square = np.ones((3, 3), dtype=bool)
    reach = max(1, int(math.ceil(min_gap_cells)))
    solve = ndimage.binary_dilation(far, square, iterations=reach) & gap
    free = solve | (ndimage.binary_dilation(solve, square, iterations=reach) & ~inside & ~known)
    # 라플라스(막 펴기)는 홀로 남은 땅 점마다 고깔이 생겨 튐이 늘어 버림(용화 실측)
    z[free] = inpaint_biharmonic(z.astype(np.float64), free)[free]
    return far


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]
    # 외곽을 그릴 지면점(다듬기 판 = 바탕 필터 마스크) — 다듬기가 지운 자리는 외곽 안 구멍으로 메움
    outer_mask: np.ndarray | None = None
    _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
    _filled: dict[float, np.ndarray] = field(default_factory=dict)

    @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 grid_sample(self, count: int, seed: int = IMPLICIT_SAMPLE_SEED) -> np.ndarray:
        """implicit 전용 — 격자 칸(√(외곽 넓이 / 점 수))마다 하나 · 모자라면 무작위로 채움.

        파일 차례 뽑기는 점 많은 트인 곳에 몰려 숲 밑 성긴 곳이 빔(요구 558). 씨앗 고정.
        """
        indices = self.indices()
        if len(indices) <= count:
            return np.asarray(self.xyz[indices], dtype=np.float32)
        # 칸 = 이 필터 자기 외곽 넓이 기준 — 바탕 외곽(766-2)으로 넓혀도 원래 자리 중심점 그대로
        own = (
            self
            if self.outer_mask is None
            else replace(self, outer_mask=None, _footprint=None, _grids={}, _filled={})
        )
        _, _, resolution, footprint = own.footprint()
        cell = max(math.sqrt(float(footprint.sum()) * resolution * resolution / count), 0.01)
        xy = np.asarray(self.xyz[indices][:, :2], dtype=np.float64)
        rng = np.random.default_rng(seed)
        order = rng.permutation(len(indices))
        gx = ((xy[order, 0] - self.bounds[0, 0]) // cell).astype(np.int64)
        gy = ((xy[order, 1] - self.bounds[1, 0]) // cell).astype(np.int64)
        _, first = np.unique(gy * (int(gx.max()) + 1) + gx, return_index=True)
        chosen = order[first]
        if len(chosen) > count:
            chosen = rng.choice(chosen, count, replace=False)
        elif len(chosen) < count:
            rest = np.setdiff1d(np.arange(len(indices)), chosen, assume_unique=True)
            chosen = np.concatenate([chosen, rng.choice(rest, count - len(chosen), replace=False)])
        return np.asarray(self.xyz[indices[np.sort(chosen)]], dtype=np.float32)

    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() if self.outer_mask is None else np.flatnonzero(self.outer_mask)
        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)
        inside = self.contains_xy(*np.meshgrid(x_coords, y_coords)).reshape(rows, cols)
        gap_cells = float(self.config.get("footprint_gap_close_meters", 1.0)) / resolution
        self._filled[resolution] = fill_holes_smooth(grid, ~missing, inside, gap_cells)
        result = (x_coords, y_coords, grid)
        self._grids[resolution] = result
        return result

    def filled(self, resolution: float) -> np.ndarray:
        """`grid(resolution)` 가 막 펴기로 메운 칸(766-2)."""
        self.grid(resolution)
        return self._filled[max(float(resolution), 0.05)]

    def hole_guides(self) -> np.ndarray:
        """implicit 이 구멍도 같은 면으로 잇게 — 메운 dtm 칸을 틈 닫기 간격마다 하나씩 (N, 3)."""
        resolution = float(self.config["dtm_grid_resolution_meters"])
        x_coords, y_coords, z_grid = self.grid(resolution)
        filled = self.filled(resolution)
        spacing = float(self.config.get("footprint_gap_close_meters", 1.0))
        step = max(1, int(round(spacing / resolution)))
        lattice = np.zeros_like(filled)
        lattice[::step, ::step] = True
        rows, cols = np.nonzero(filled & lattice)
        return np.column_stack([x_coords[cols], y_coords[rows], z_grid[rows, cols]]).astype(
            np.float32
        )

    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._filled.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