"""B04 LAS/LAZ 고속 구조화 엔진 — 여러 장을 한 벌로 병합한다 (2026-09-06 사용자 확정).

드론 라이다는 사업지가 넓으면 도엽별로 여러 장이 온다. 여기서 **합친 범위**로 한 벌을
만들고, 뒤 단계(지면필터·모델·등고선·배수)는 받는 형식이 그대로라 손대지 않는다.

점이 임계를 넘으면 **칸(기본 0.5m)마다 최저점 하나만** 남긴다(씨닝). 설계가 쓰는 격자가
1m(지면필터 2m·CSF 천 1.5m)라 0.5m 는 설계보다 촘촘해 결과 표고가 사실상 같고, 30GB 두
장이 메모리 29GB → 1.4GB 로 내려간다. 임계 아래면 원본 점을 그대로 쓴다 — 작은 자료의
결과는 바뀌지 않는다.
"""

import logging
import os
import tempfile
from collections.abc import Callable, Sequence
from pathlib import Path
from typing import Any

import laspy
import numpy as np

from common_util.common_util_json import replace_with_retry
from config.config_system import (
    SURFACE_DEFAULT_RGB_VALUE,
    SURFACE_LAS_CHUNK_SIZE,
    SURFACE_MERGE_MAX_GAP_M,
    SURFACE_THIN_CELL_SIZE_M,
    SURFACE_THIN_MAX_CELLS,
    SURFACE_THIN_TRIGGER_POINTS,
)

logger = logging.getLogger(__name__)

ProgressCallback = Callable[[int], None]
PathLike = str | Path
# 청크에서 점과 함께 옮기는 속성들 — 파일에 없으면 기본값이 남는다.
_ATTRIBUTES = ("intensity", "rgb", "return_number", "number_of_returns", "classification")


def _as_list(las_path: PathLike | Sequence[PathLike]) -> list[Path]:
    if isinstance(las_path, (str, Path)):
        return [Path(las_path)]
    return [Path(item) for item in las_path]


def point_cloud_extent(path: PathLike) -> tuple[int, tuple[float, float, float, float]]:
    """머리글만 읽어 점 수와 XY 범위를 돌려준다 — 파일 크기와 무관하게 즉시 끝난다."""
    with laspy.open(Path(path)) as las_file:
        header = las_file.header
        return int(header.point_count), (
            float(header.mins[0]),
            float(header.mins[1]),
            float(header.maxs[0]),
            float(header.maxs[1]),
        )


def merge_gap_error(
    paths: Sequence[PathLike], gap_m: float = SURFACE_MERGE_MAX_GAP_M
) -> str | None:
    """서로 멀리 떨어진 지형 파일이 섞였는지 — 문제면 안내 문구, 없으면 None.

    다른 사업지 파일이나 좌표계가 다른 파일이 섞이면 합친 범위가 통째로 어긋나 격자가
    터진다. 도엽으로 나뉜 자료는 경계가 맞닿으므로 여유를 두고 **어느 파일과도 만나지
    않는 파일**만 걸러 낸다.
    """
    sources = _as_list(paths)
    if len(sources) < 2:
        return None
    boxes = [(path, point_cloud_extent(path)[1]) for path in sources]
    for index, (path, box) in enumerate(boxes):
        near = any(
            box[0] - gap_m <= other[2]
            and other[0] - gap_m <= box[2]
            and box[1] - gap_m <= other[3]
            and other[1] - gap_m <= box[3]
            for other_index, (_, other) in enumerate(boxes)
            if other_index != index
        )
        if not near:
            return (
                f"지형 파일 「{path.name}」의 좌표가 다른 파일과 "
                f"{int(gap_m):,}m 넘게 떨어져 있습니다."
                " 같은 사업지의 파일인지, 좌표계가 같은지 확인해 주십시오."
            )
    return None


class _Merged:
    """합친 점을 담는 그릇 — 원본 유지형과 씨닝형이 같은 모양으로 낸다."""

    def __init__(self, capacity: int) -> None:
        self.xyz = np.empty((capacity, 3), dtype=np.float64)
        self.intensity = np.zeros(capacity, dtype=np.uint16)
        self.rgb = np.full((capacity, 3), SURFACE_DEFAULT_RGB_VALUE, dtype=np.uint8)
        self.return_number = np.ones(capacity, dtype=np.uint8)
        self.number_of_returns = np.ones(capacity, dtype=np.uint8)
        self.classification = np.zeros(capacity, dtype=np.uint8)
        self.size = 0

    def arrays(self) -> dict[str, np.ndarray]:
        end = self.size
        return {
            "xyz": self.xyz[:end],
            "intensity": self.intensity[:end],
            "rgb": self.rgb[:end],
            "return_number": self.return_number[:end],
            "number_of_returns": self.number_of_returns[:end],
            "classification": self.classification[:end],
        }

    def append(self, columns: dict[str, np.ndarray]) -> None:
        count = len(columns["x"])
        section = slice(self.size, self.size + count)
        self.xyz[section, 0] = columns["x"]
        self.xyz[section, 1] = columns["y"]
        self.xyz[section, 2] = columns["z"]
        for key in _ATTRIBUTES:
            if key in columns:
                getattr(self, key)[section] = columns[key]
        self.size += count


def _chunk_columns(chunk: Any, dimensions: set[str]) -> dict[str, np.ndarray]:
    """청크에서 쓸 값만 꺼낸다. 파일에 없는 항목은 키를 빼서 기본값이 남게 한다."""
    columns: dict[str, np.ndarray] = {
        "x": np.asarray(chunk.x, dtype=np.float64),
        "y": np.asarray(chunk.y, dtype=np.float64),
        "z": np.asarray(chunk.z, dtype=np.float64),
    }
    if "intensity" in dimensions:
        columns["intensity"] = np.asarray(chunk.intensity, dtype=np.uint16)
    if {"red", "green", "blue"}.issubset(dimensions):
        colors = np.stack(
            [
                np.asarray(chunk.red, dtype=np.float64),
                np.asarray(chunk.green, dtype=np.float64),
                np.asarray(chunk.blue, dtype=np.float64),
            ],
            axis=1,
        )
        if colors.size and float(colors.max()) > 255.0:
            colors /= 256.0
        columns["rgb"] = colors.clip(0, 255).astype(np.uint8)
    if {"return_number", "number_of_returns"}.issubset(dimensions):
        columns["return_number"] = np.asarray(chunk.return_number, dtype=np.uint8)
        columns["number_of_returns"] = np.asarray(chunk.number_of_returns, dtype=np.uint8)
    if "classification" in dimensions:
        columns["classification"] = np.asarray(chunk.classification, dtype=np.uint8)
    return columns


class _ThinGrid:
    """씨닝형 — **지면 분류점은 전부** 남기고, 나머지는 칸마다 최저점 하나만 남긴다.

    설계 지표면을 만드는 것은 지면점이다(업체가 분류해 준 ASPRS class 2). 그 점을 하나도
    버리지 않으므로 **지면 결과는 씨닝 전과 완전히 같다**(2026-09-06 용화 실측: 1m 지면
    격자 141,969칸 전부 표고 차이 0). 지면점은 원본의 2~3%뿐이라 남겨도 가볍다.

    나머지(수목·구조물·잡음)는 칸마다 최저점만 남긴다 — 분류가 없는 자료에서 CSF·PMF가
    지면을 찾을 밑그림으로 충분하다(CSF 천 간격 1.5m > 칸 0.5m).

    한 청크 안에서 같은 칸이 여러 번 나오면 뒤에 쓴 값이 이겨 최저점이 아니게 된다.
    그래서 청크를 (칸, 표고)로 정렬해 **칸마다 첫 점**만 골라 낸 뒤 격자와 견준다.
    """

    #: ASPRS 지면 분류 코드.
    GROUND_CLASS = 2

    def __init__(self, bounds: np.ndarray, cell_size: float) -> None:
        self.cell_size = cell_size
        self.x_min = float(bounds[0, 0])
        self.y_min = float(bounds[1, 0])
        self.width = int(np.ceil((float(bounds[0, 1]) - self.x_min) / cell_size)) + 1
        self.height = int(np.ceil((float(bounds[1, 1]) - self.y_min) / cell_size)) + 1
        cells = self.width * self.height
        if cells > SURFACE_THIN_MAX_CELLS:
            raise ValueError(
                "지형 자료의 합친 범위가 너무 넓습니다."
                " 같은 사업지의 파일인지, 좌표계가 같은지 확인해 주십시오."
            )
        self.best_z = np.full(cells, np.inf, dtype=np.float64)
        self.x = np.zeros(cells, dtype=np.float64)
        self.y = np.zeros(cells, dtype=np.float64)
        self.intensity = np.zeros(cells, dtype=np.uint16)
        self.rgb = np.full((cells, 3), SURFACE_DEFAULT_RGB_VALUE, dtype=np.uint8)
        self.return_number = np.ones(cells, dtype=np.uint8)
        self.number_of_returns = np.ones(cells, dtype=np.uint8)
        self.classification = np.zeros(cells, dtype=np.uint8)
        # 그대로 남길 지면점 — 청크마다 모아 두었다가 마지막에 잇는다.
        self.ground: list[dict[str, np.ndarray]] = []

    def add(self, columns: dict[str, np.ndarray]) -> None:
        classification = columns.get("classification")
        if classification is not None:
            is_ground = classification == self.GROUND_CLASS
            if is_ground.any():
                self.ground.append({key: value[is_ground] for key, value in columns.items()})
                keep = ~is_ground
                columns = {key: value[keep] for key, value in columns.items()}
        x, y, z = columns["x"], columns["y"], columns["z"]
        if not len(x):
            return
        grid_x = np.clip(((x - self.x_min) / self.cell_size).astype(np.int64), 0, self.width - 1)
        grid_y = np.clip(((y - self.y_min) / self.cell_size).astype(np.int64), 0, self.height - 1)
        cell = grid_y * self.width + grid_x
        order = np.lexsort((z, cell))
        sorted_cell = cell[order]
        first = np.ones(len(order), dtype=bool)
        first[1:] = sorted_cell[1:] != sorted_cell[:-1]
        candidate = order[first]
        candidate_cell = cell[candidate]
        better = z[candidate] < self.best_z[candidate_cell]
        chosen = candidate[better]
        target = candidate_cell[better]
        self.best_z[target] = z[chosen]
        self.x[target] = x[chosen]
        self.y[target] = y[chosen]
        for key in _ATTRIBUTES:
            if key in columns:
                getattr(self, key)[target] = columns[key][chosen]

    def collect(self) -> _Merged:
        occupied = np.flatnonzero(np.isfinite(self.best_z))
        ground_count = sum(len(item["x"]) for item in self.ground)
        merged = _Merged(len(occupied) + ground_count)
        merged.xyz[: len(occupied), 0] = self.x[occupied]
        merged.xyz[: len(occupied), 1] = self.y[occupied]
        merged.xyz[: len(occupied), 2] = self.best_z[occupied]
        for key in _ATTRIBUTES:
            getattr(merged, key)[: len(occupied)] = getattr(self, key)[occupied]
        merged.size = len(occupied)
        for item in self.ground:
            merged.append(item)
        return merged


def _headers(sources: list[Path]) -> tuple[int, np.ndarray, bool]:
    """전체 점 수·합친 범위(3x2)·색 보유 여부를 머리글만 읽어 구한다."""
    total = 0
    has_rgb = False
    mins = np.full(3, np.inf, dtype=np.float64)
    maxs = np.full(3, -np.inf, dtype=np.float64)
    for source in sources:
        with laspy.open(source) as las_file:
            header = las_file.header
            total += int(header.point_count)
            mins = np.minimum(mins, np.asarray(header.mins, dtype=np.float64))
            maxs = np.maximum(maxs, np.asarray(header.maxs, dtype=np.float64))
            dimensions = set(header.point_format.dimension_names)
            has_rgb = has_rgb or {"red", "green", "blue"}.issubset(dimensions)
    if not np.isfinite(mins).all():
        mins = np.zeros(3, dtype=np.float64)
        maxs = np.zeros(3, dtype=np.float64)
    return total, np.column_stack((mins, maxs)), has_rgb


def _merge_sources(
    sources: list[Path],
    total_points: int,
    bounds: np.ndarray,
    thin: bool,
    progress_callback: ProgressCallback | None,
) -> _Merged:
    grid = _ThinGrid(bounds, SURFACE_THIN_CELL_SIZE_M) if thin else None
    merged = _Merged(total_points) if grid is None else None
    done = 0
    for source in sources:
        with laspy.open(source) as las_file:
            dimensions = set(las_file.header.point_format.dimension_names)
            for chunk in las_file.chunk_iterator(SURFACE_LAS_CHUNK_SIZE):
                columns = _chunk_columns(chunk, dimensions)
                if grid is not None:
                    grid.add(columns)
                else:
                    merged.append(columns)
                done += len(columns["x"])
                if progress_callback:
                    progress_callback(int(done / total_points * 100) if total_points else 100)
    return grid.collect() if grid is not None else merged


def structurize_las(
    las_path: PathLike | Sequence[PathLike],
    output_dir: str | Path,
    progress_callback: ProgressCallback | None = None,
) -> Path:
    """지형 파일 한 장 또는 여러 장을 청크로 읽어 B04 structured.npz로 원자적 저장한다."""
    sources = _as_list(las_path)
    if not sources:
        raise ValueError("구조화할 지형 파일이 없습니다.")
    target_dir = Path(output_dir)
    target_dir.mkdir(parents=True, exist_ok=True)
    target = target_dir / "structured.npz"

    total_points, bounds, has_rgb = _headers(sources)
    thin = total_points > SURFACE_THIN_TRIGGER_POINTS
    merged = _merge_sources(sources, total_points, bounds, thin, progress_callback)
    logger.info(
        "B04 구조화: 파일 %d장 원본 %d점 → 저장 %d점 (씨닝 %s)",
        len(sources),
        total_points,
        merged.size,
        f"{SURFACE_THIN_CELL_SIZE_M}m 칸" if thin else "없음",
    )

    temporary_path: Path | None = None
    try:
        with tempfile.NamedTemporaryFile(
            mode="wb",
            dir=target_dir,
            prefix=".structured.",
            suffix=".npz.tmp",
            delete=False,
        ) as temporary:
            temporary_path = Path(temporary.name)
            np.savez_compressed(
                temporary,
                **merged.arrays(),
                bounds=bounds,
                total_points=np.array([merged.size], dtype=np.int64),
                source_point_count=np.array([total_points], dtype=np.int64),
                source_file_count=np.array([len(sources)], dtype=np.int64),
                thinned=np.array([int(thin)], dtype=np.int8),
                has_rgb=np.array([int(has_rgb)], dtype=np.int8),
            )
            temporary.flush()
            os.fsync(temporary.fileno())
        # 이 교체가 윈도우 공유 위반으로 죽으면 WF1 분석 전체가 실패한다 — 재시도한다.
        replace_with_retry(temporary_path, target)
        temporary_path = None
    finally:
        if temporary_path is not None:
            temporary_path.unlink(missing_ok=True)

    if progress_callback:
        progress_callback(100)
    return target