"""Provider-neutral, dependency-light LiDAR ground segmentation primitives."""

from __future__ import annotations

import hashlib
import importlib
import math
import platform
import time
from collections.abc import Mapping
from dataclasses import dataclass
from pathlib import Path
from types import ModuleType
from typing import Final, Protocol

import numpy as np
import numpy.typing as npt

BoolArray = npt.NDArray[np.bool_]
PATCHWORKPP_SOURCE_URL: Final = "https://github.com/url-kaist/patchwork-plusplus"
PATCHWORKPP_SOURCE_TAG: Final = "v1.4.1"
PATCHWORKPP_SOURCE_COMMIT: Final = "3e6903a1d5537a4cc2ace897b0bbb98a92d6014c"


class GroundSegmentationError(ValueError):
    """A ground provider or profile violates the shared segmentation contract."""


@dataclass(frozen=True, slots=True)
class GroundBenchmarkProfile:
    profile_id: str = "k1-vendor-map-ground-ab/v1"
    pose_binding_threshold_ms: float = 100.0
    current_cell_size_m: float = 0.5
    current_local_radius_m: float = 2.5
    current_lower_percentile: float = 8.0
    current_maximum_below_ground_m: float = 0.25
    current_maximum_above_ground_m: float = 0.12
    current_minimum_local_points: int = 8
    patchwork_sensor_height_proxy_m: float = 0.0
    patchwork_map_vertical_origin_offset_m: float = 0.0
    patchwork_height_evidence: str = "missing"
    patchwork_minimum_range_m: float = 0.1
    patchwork_maximum_range_m: float = 20.0

    def __post_init__(self) -> None:
        finite_values = (
            self.pose_binding_threshold_ms,
            self.current_cell_size_m,
            self.current_local_radius_m,
            self.current_lower_percentile,
            self.current_maximum_below_ground_m,
            self.current_maximum_above_ground_m,
            self.patchwork_sensor_height_proxy_m,
            self.patchwork_map_vertical_origin_offset_m,
            self.patchwork_minimum_range_m,
            self.patchwork_maximum_range_m,
        )
        if not all(math.isfinite(value) for value in finite_values):
            raise GroundSegmentationError("Ground benchmark profile must be finite")
        if (
            not 0 < self.pose_binding_threshold_ms <= 10_000
            or not 0 < self.current_cell_size_m <= 100
            or not 0 < self.current_local_radius_m <= 1_000
            or not 0 <= self.current_lower_percentile <= 100
            or not 0 <= self.current_maximum_below_ground_m <= 100
            or not 0 <= self.current_maximum_above_ground_m <= 100
            or not 1 <= self.current_minimum_local_points <= 1_000_000
            or not 0 <= self.patchwork_sensor_height_proxy_m <= 10
            or not -10 <= self.patchwork_map_vertical_origin_offset_m <= 10
            or not 0 <= self.patchwork_minimum_range_m < self.patchwork_maximum_range_m <= 1_000
            or self.patchwork_height_evidence
            not in {"missing", "operator-estimated", "runtime-calibrated"}
        ):
            raise GroundSegmentationError("Ground benchmark profile is invalid")
        if self.patchwork_height_evidence == "missing" and (
            self.patchwork_sensor_height_proxy_m != 0
            or self.patchwork_map_vertical_origin_offset_m != 0
        ):
            raise GroundSegmentationError(
                "Ground benchmark cannot apply height without height evidence"
            )
        if (
            self.patchwork_height_evidence != "missing"
            and self.patchwork_sensor_height_proxy_m <= 0
        ):
            raise GroundSegmentationError(
                "Ground benchmark height evidence requires a positive sensor height"
            )

    def to_dict(self) -> dict[str, object]:
        return {
            "schema_version": "missioncore.lidar-ground-benchmark-profile/v1",
            "profile_id": self.profile_id,
            "pose_binding": {
                "basis": "nearest-recorded-host-monotonic-arrival",
                "threshold_ms": self.pose_binding_threshold_ms,
            },
            "current_baseline": {
                "provider_id": "missioncore-local-percentile-ground/v1",
                "derived_from": "local-ground-relative-object-support-v1",
                "cell_size_m": self.current_cell_size_m,
                "local_radius_m": self.current_local_radius_m,
                "lower_percentile": self.current_lower_percentile,
                "maximum_below_ground_m": self.current_maximum_below_ground_m,
                "maximum_above_ground_m": self.current_maximum_above_ground_m,
                "minimum_local_points": self.current_minimum_local_points,
            },
            "candidate": {
                "provider_id": "patchworkpp/v1.4.1",
                "sensor_height_proxy_m": self.patchwork_sensor_height_proxy_m,
                "map_vertical_origin_offset_m": self.patchwork_map_vertical_origin_offset_m,
                "height_evidence": self.patchwork_height_evidence,
                "minimum_range_m": self.patchwork_minimum_range_m,
                "maximum_range_m": self.patchwork_maximum_range_m,
                "enable_rnr": True,
                "enable_rvpf": True,
                "enable_tgr": True,
            },
            "input_normalization": {
                "current": "vendor-map-xyz",
                "candidate": "best-effort-map-to-lidar-pose-inversion",
                "physical_sensor_height_known": (
                    self.patchwork_height_evidence == "runtime-calibrated"
                ),
                "sensor_scan_geometry_known": False,
            },
            "authority": {
                "commands_enabled": False,
                "navigation_or_safety_accepted": False,
            },
        }


DEFAULT_GROUND_BENCHMARK_PROFILE: Final = GroundBenchmarkProfile()


@dataclass(frozen=True, slots=True)
class GroundSegmentation:
    ground_mask: BoolArray
    assigned_mask: BoolArray
    latency_ms: float


class GroundSegmenter(Protocol):
    @property
    def identity(self) -> Mapping[str, object]: ...

    def segment(self, xyzi: npt.NDArray[np.float32]) -> GroundSegmentation: ...


class PatchworkGroundProfile(Protocol):
    @property
    def patchwork_sensor_height_proxy_m(self) -> float: ...

    @property
    def patchwork_minimum_range_m(self) -> float: ...

    @property
    def patchwork_maximum_range_m(self) -> float: ...


class PatchworkPPGroundSegmenter:
    """Runtime-only adapter for the pinned official Patchwork++ Python binding."""

    def __init__(
        self,
        module: ModuleType,
        profile: PatchworkGroundProfile,
        *,
        source_commit: str = PATCHWORKPP_SOURCE_COMMIT,
        source_tag: str = PATCHWORKPP_SOURCE_TAG,
    ) -> None:
        if len(source_commit) != 40 or any(
            character not in "0123456789abcdef" for character in source_commit
        ):
            raise GroundSegmentationError("Patchwork++ source commit is invalid")
        if source_tag != PATCHWORKPP_SOURCE_TAG:
            raise GroundSegmentationError("Patchwork++ source tag is not admitted")
        module_path_value = getattr(module, "__file__", None)
        if not isinstance(module_path_value, str):
            raise GroundSegmentationError("Patchwork++ module has no verifiable binary")
        module_path = Path(module_path_value).resolve(strict=True)
        params = module.Parameters()
        params.sensor_height = profile.patchwork_sensor_height_proxy_m
        params.min_range = profile.patchwork_minimum_range_m
        params.max_range = profile.patchwork_maximum_range_m
        params.enable_RNR = True
        params.enable_RVPF = True
        params.enable_TGR = True
        params.verbose = False
        self._estimator = module.patchworkpp(params)
        self._identity = {
            "provider_id": "patchworkpp/v1.4.1",
            "source_url": PATCHWORKPP_SOURCE_URL,
            "source_tag": source_tag,
            "source_commit": source_commit,
            "binding_version": str(getattr(module, "__version__", "unknown")),
            "binary_sha256": _sha256(module_path),
            "platform": platform.system().lower(),
            "machine": platform.machine().lower(),
            "ground_truth": False,
        }

    @classmethod
    def load(
        cls,
        profile: PatchworkGroundProfile = DEFAULT_GROUND_BENCHMARK_PROFILE,
        *,
        module_name: str = "pypatchworkpp",
        source_commit: str = PATCHWORKPP_SOURCE_COMMIT,
        source_tag: str = PATCHWORKPP_SOURCE_TAG,
    ) -> PatchworkPPGroundSegmenter:
        try:
            module = importlib.import_module(module_name)
        except ImportError as exc:
            raise GroundSegmentationError(
                "Pinned Patchwork++ Python binding is unavailable"
            ) from exc
        return cls(
            module,
            profile,
            source_commit=source_commit,
            source_tag=source_tag,
        )

    @property
    def identity(self) -> Mapping[str, object]:
        return self._identity

    def segment(self, xyzi: npt.NDArray[np.float32]) -> GroundSegmentation:
        points = np.ascontiguousarray(_xyzi(xyzi), dtype=np.float32)
        started = time.perf_counter_ns()
        self._estimator.estimateGround(points)
        latency_ms = (time.perf_counter_ns() - started) / 1_000_000
        ground_indices = np.asarray(
            self._estimator.getGroundIndices(),
            dtype=np.int64,
        ).reshape((-1,))
        nonground_indices = np.asarray(
            self._estimator.getNongroundIndices(),
            dtype=np.int64,
        ).reshape((-1,))
        _indices(ground_indices, points.shape[0], "Patchwork++ ground")
        _indices(nonground_indices, points.shape[0], "Patchwork++ non-ground")
        if np.intersect1d(ground_indices, nonground_indices).size:
            raise GroundSegmentationError("Patchwork++ assigned one point twice")
        ground = np.zeros(points.shape[0], dtype=np.bool_)
        assigned = np.zeros(points.shape[0], dtype=np.bool_)
        ground[ground_indices] = True
        assigned[ground_indices] = True
        assigned[nonground_indices] = True
        return GroundSegmentation(
            ground_mask=ground,
            assigned_mask=assigned,
            latency_ms=latency_ms,
        )


class LocalPercentileGroundSegmenter:
    """Full-frame diagnostic extension of the existing E19 local ground heuristic."""

    def __init__(self, profile: GroundBenchmarkProfile) -> None:
        self.profile = profile

    @property
    def identity(self) -> Mapping[str, object]:
        return {
            "provider_id": "missioncore-local-percentile-ground/v1",
            "implementation_sha256": _sha256(Path(__file__).resolve(strict=True)),
            "ground_truth": False,
        }

    def segment(self, xyzi: npt.NDArray[np.float32]) -> GroundSegmentation:
        points = _xyzi(xyzi)
        started = time.perf_counter_ns()
        xyz = points[:, :3].astype(np.float64, copy=False)
        cell_size = self.profile.current_cell_size_m
        cell_keys = np.floor(xyz[:, :2] / cell_size).astype(np.int64)
        unique_cells, inverse = np.unique(cell_keys, axis=0, return_inverse=True)
        order = np.argsort(inverse, kind="stable")
        counts = np.bincount(inverse, minlength=unique_cells.shape[0])
        offsets = np.concatenate(([0], np.cumsum(counts)))
        cell_lookup = {
            (int(cell[0]), int(cell[1])): cell_index for cell_index, cell in enumerate(unique_cells)
        }
        neighbor_span = math.ceil(self.profile.current_local_radius_m / cell_size) + 1
        ground = np.zeros(points.shape[0], dtype=np.bool_)
        global_ground_z = float(np.percentile(xyz[:, 2], self.profile.current_lower_percentile))
        radius_squared = self.profile.current_local_radius_m**2
        for cell_index in range(unique_cells.shape[0]):
            point_indices = order[offsets[cell_index] : offsets[cell_index + 1]]
            if point_indices.size == 0:
                continue
            center_xy = np.median(xyz[point_indices, :2], axis=0)
            cell_x, cell_y = unique_cells[cell_index]
            neighbor_slices: list[npt.NDArray[np.int64]] = []
            for delta_x in range(-neighbor_span, neighbor_span + 1):
                for delta_y in range(-neighbor_span, neighbor_span + 1):
                    neighbor_cell_index = cell_lookup.get(
                        (int(cell_x + delta_x), int(cell_y + delta_y))
                    )
                    if neighbor_cell_index is None:
                        continue
                    neighbor_slices.append(
                        order[offsets[neighbor_cell_index] : offsets[neighbor_cell_index + 1]]
                    )
            local_indices = np.concatenate(neighbor_slices)
            local_xyz = xyz[local_indices]
            delta_xy = local_xyz[:, :2] - center_xy
            local = local_xyz[
                np.einsum("ij,ij->i", delta_xy, delta_xy) <= radius_squared,
                2,
            ]
            ground_z = (
                float(np.percentile(local, self.profile.current_lower_percentile))
                if local.size >= self.profile.current_minimum_local_points
                else global_ground_z
            )
            z = xyz[point_indices, 2]
            ground[point_indices] = (
                z >= ground_z - self.profile.current_maximum_below_ground_m
            ) & (z <= ground_z + self.profile.current_maximum_above_ground_m)
        latency_ms = (time.perf_counter_ns() - started) / 1_000_000
        return GroundSegmentation(
            ground_mask=ground,
            assigned_mask=np.ones(points.shape[0], dtype=np.bool_),
            latency_ms=latency_ms,
        )


def _xyzi(value: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
    points = np.asarray(value)
    if (
        points.dtype != np.dtype(np.float32)
        or points.ndim != 2
        or points.shape[1] != 4
        or points.shape[0] == 0
        or not np.isfinite(points).all()
    ):
        raise GroundSegmentationError("Ground input must be a non-empty finite float32 XYZI matrix")
    return points


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1024**2), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _indices(value: npt.NDArray[np.int64], count: int, label: str) -> None:
    if value.ndim != 1 or np.any(value < 0) or np.any(value >= count):
        raise GroundSegmentationError(f"{label} indices are invalid")
    if np.unique(value).size != value.size:
        raise GroundSegmentationError(f"{label} indices are duplicated")
