from __future__ import annotations

import json
import os
import tempfile
import threading
import time
from collections import OrderedDict
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from urllib.parse import urlparse
from urllib.request import urlopen

import cv2
import numpy as np

try:
    import onnxruntime as ort
except ImportError:  # pragma: no cover
    ort = None

try:
    import open3d as o3d
except ImportError:  # pragma: no cover
    o3d = None

try:
    from minio import Minio
except ImportError:  # pragma: no cover
    Minio = None


@dataclass(frozen=True)
class CpuSettings:
    profile: str
    registry_path: Path
    model_dir: Path
    execution_provider: str
    intra_op_threads: int
    inter_op_threads: int
    max_concurrency: int
    max_loaded_models: int
    max_resource_bytes: int
    max_point_count: int
    fallback_mode: str

    @classmethod
    def from_env(cls) -> "CpuSettings":
        return cls(
            profile=os.getenv("RAIL_RUNTIME_PROFILE", "cpu-local"),
            registry_path=Path(os.getenv("RAIL_MODEL_REGISTRY", "/app/config/model-registry.json")),
            model_dir=Path(os.getenv("RAIL_MODEL_DIR", "/models")),
            execution_provider=os.getenv("RAIL_EXECUTION_PROVIDER", "CPUExecutionProvider"),
            intra_op_threads=max(1, int(os.getenv("RAIL_INTRA_OP_THREADS", "4"))),
            inter_op_threads=max(1, int(os.getenv("RAIL_INTER_OP_THREADS", "1"))),
            max_concurrency=max(1, int(os.getenv("RAIL_MAX_CONCURRENCY", "1"))),
            max_loaded_models=max(1, int(os.getenv("RAIL_MAX_LOADED_MODELS", "1"))),
            max_resource_bytes=max(1, int(os.getenv("RAIL_MAX_RESOURCE_MB", "256"))) * 1024 * 1024,
            max_point_count=max(1000, int(os.getenv("RAIL_MAX_POINT_COUNT", "500000"))),
            fallback_mode=os.getenv("RAIL_FALLBACK_MODE", "baseline"),
        )


@dataclass(frozen=True)
class SpatialModelSpec:
    model_group: str
    model_version: str
    display_name: str
    family: str
    runtime: str
    artifact: str | None
    parser: str
    input_size: int | None
    precision: str
    mean: tuple[float, float, float]
    std: tuple[float, float, float]
    active: bool

    @classmethod
    def from_dict(cls, value: dict[str, Any]) -> "SpatialModelSpec":
        return cls(
            model_group=str(value["model_group"]),
            model_version=str(value.get("model_version", "cpu-v1.0.0")),
            display_name=str(value.get("display_name", value["model_group"])),
            family=str(value.get("family", "custom")),
            runtime=str(value.get("runtime", "onnxruntime-cpu")),
            artifact=value.get("artifact"),
            parser=str(value.get("parser", "binary-change")),
            input_size=int(value["input_size"]) if value.get("input_size") else None,
            precision=str(value.get("precision", "FP32")),
            mean=tuple(float(item) for item in value.get("mean", [0.485, 0.456, 0.406])),
            std=tuple(float(item) for item in value.get("std", [0.229, 0.224, 0.225])),
            active=bool(value.get("active", False)),
        )

    def artifact_path(self, model_dir: Path) -> Path | None:
        return model_dir / self.artifact if self.artifact else None


@dataclass(frozen=True)
class SpatialOutcome:
    execution_mode: str
    results: list[dict[str, Any]]
    reason: str | None
    latency_ms: float
    model: dict[str, Any]


class SpatialRegistry:
    def __init__(self, settings: CpuSettings):
        self.settings = settings
        payload = self._payload()
        self.profile = str(payload.get("profile", settings.profile))
        self._specs: dict[str, list[SpatialModelSpec]] = {}
        for spec in (SpatialModelSpec.from_dict(item) for item in payload.get("models", [])):
            if spec.model_group in {"change-detector", "pointcloud-analyzer"}:
                self._specs.setdefault(spec.model_group, []).append(spec)

    def get(self, model_group: str, model_version: str | None = None) -> SpatialModelSpec | None:
        specs = self._specs.get(model_group, [])
        if model_version:
            match = next((item for item in specs if item.model_version == model_version), None)
            return match
        return next((item for item in specs if item.active), specs[0] if specs else None)

    def all(self) -> list[SpatialModelSpec]:
        return [spec for specs in self._specs.values() for spec in specs]

    def _payload(self) -> dict[str, Any]:
        if self.settings.registry_path.is_file():
            return json.loads(self.settings.registry_path.read_text(encoding="utf-8"))
        return {
            "profile": self.settings.profile,
            "models": [
                {
                    "model_group": "change-detector",
                    "model_version": "cpu-v1.0.0",
                    "display_name": "BIT-R18 CPU",
                    "family": "BIT",
                    "runtime": "onnxruntime-cpu",
                    "artifact": "change-detector/bit-r18/model.onnx",
                    "parser": "binary-change",
                    "input_size": 512,
                },
                {
                    "model_group": "pointcloud-analyzer",
                    "model_version": "cpu-v1.0.0",
                    "display_name": "Open3D 几何分析 CPU",
                    "family": "Open3D",
                    "runtime": "open3d-cpu",
                    "artifact": None,
                    "parser": "pointcloud-geometry",
                    "precision": "FP64",
                },
            ],
        }


class ObjectReader:
    def __init__(self, settings: CpuSettings):
        self.settings = settings
        self._minio: Any = None

    def read(self, location: str) -> bytes:
        parsed = urlparse(location)
        if parsed.scheme == "s3":
            return self._read_s3(parsed.netloc, parsed.path.lstrip("/"))
        if parsed.scheme == "file":
            return self._bounded(Path(parsed.path).read_bytes())
        if parsed.scheme in {"http", "https"}:
            with urlopen(location, timeout=30) as response:  # noqa: S310 - controlled deployment input
                return self._bounded(response.read(self.settings.max_resource_bytes + 1))
        path = Path(location)
        if path.is_file():
            return self._bounded(path.read_bytes())
        raise FileNotFoundError(f"资源不存在或协议不受支持: {location}")

    def image(self, location: str) -> np.ndarray:
        image = cv2.imdecode(np.frombuffer(self.read(location), dtype=np.uint8), cv2.IMREAD_COLOR)
        if image is None:
            raise ValueError("资源不是可解码图像")
        return image

    def point_cloud(self, location: str) -> Any:
        if o3d is None:
            raise RuntimeError("open3d 未安装")
        suffix = Path(urlparse(location).path).suffix.lower() or ".ply"
        if suffix not in {".ply", ".pcd", ".xyz", ".xyzn", ".xyzrgb", ".pts"}:
            raise ValueError(f"CPU 模式暂不支持点云格式 {suffix}，请转换为 PLY/PCD/XYZ/PTS")
        payload = self.read(location)
        temp_path = ""
        try:
            with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as stream:
                stream.write(payload)
                temp_path = stream.name
            cloud = o3d.io.read_point_cloud(temp_path)
            if cloud.is_empty():
                raise ValueError("点云为空或文件无法解析")
            return cloud
        finally:
            if temp_path:
                Path(temp_path).unlink(missing_ok=True)

    def _read_s3(self, bucket: str, object_name: str) -> bytes:
        if Minio is None:
            raise RuntimeError("未安装 MinIO Python SDK")
        if self._minio is None:
            endpoint = os.getenv("MINIO_ENDPOINT", "minio:9000")
            secure = endpoint.startswith("https://")
            endpoint = endpoint.removeprefix("http://").removeprefix("https://")
            self._minio = Minio(
                endpoint,
                access_key=os.getenv("MINIO_ACCESS_KEY", "minioadmin"),
                secret_key=os.getenv("MINIO_SECRET_KEY", "minioadmin"),
                secure=secure,
            )
        response = self._minio.get_object(bucket, object_name)
        try:
            return self._bounded(response.read(self.settings.max_resource_bytes + 1))
        finally:
            response.close()
            response.release_conn()

    def _bounded(self, payload: bytes) -> bytes:
        if len(payload) > self.settings.max_resource_bytes:
            raise ValueError(f"资源超过 {self.settings.max_resource_bytes // (1024 * 1024)} MB CPU 模式上限")
        return payload


class SpatialCpuRuntime:
    def __init__(self):
        self.settings = CpuSettings.from_env()
        self.registry = SpatialRegistry(self.settings)
        self.reader = ObjectReader(self.settings)
        self._sessions: OrderedDict[str, Any] = OrderedDict()
        self._session_lock = threading.RLock()
        self._semaphore = threading.BoundedSemaphore(self.settings.max_concurrency)

    def analyze_change(
        self,
        storage_url: str | None,
        metadata: dict[str, Any],
        parameters: dict[str, Any],
        model_version: str | None = None,
    ) -> SpatialOutcome:
        spec = self.registry.get("change-detector", model_version)
        if spec is None:
            return self._unavailable("change-detector", "变化检测未配置")
        baseline_url = metadata.get("baseline_storage_url")
        if not storage_url or not baseline_url:
            return self._unavailable("change-detector", "真实变化检测需要当前期与 baseline_storage_url 基准影像")
        started = time.perf_counter()
        try:
            with self._semaphore:
                current = self.reader.image(storage_url)
                baseline = self.reader.image(str(baseline_url))
                artifact = spec.artifact_path(self.settings.model_dir)
                if artifact and artifact.is_file() and ort is not None:
                    results = self._onnx_change(spec, baseline, current, parameters)
                    mode = spec.runtime
                else:
                    results = self._opencv_change(baseline, current, parameters)
                    mode = "opencv-cpu"
            return SpatialOutcome(
                mode,
                results,
                None,
                round((time.perf_counter() - started) * 1000, 2),
                self._model_info(spec),
            )
        except Exception as exc:
            return SpatialOutcome(
                "baseline",
                [],
                f"变化检测失败: {type(exc).__name__}: {exc}",
                round((time.perf_counter() - started) * 1000, 2),
                self._model_info(spec),
            )

    def analyze_pointcloud(
        self,
        storage_url: str | None,
        metadata: dict[str, Any],
        parameters: dict[str, Any],
        model_version: str | None = None,
    ) -> SpatialOutcome:
        spec = self.registry.get("pointcloud-analyzer", model_version)
        if spec is None:
            return self._unavailable("pointcloud-analyzer", "点云分析未配置")
        if o3d is None:
            return self._unavailable("pointcloud-analyzer", "open3d 未安装")
        if not storage_url:
            return self._unavailable("pointcloud-analyzer", "资源没有 storage_url")
        started = time.perf_counter()
        try:
            with self._semaphore:
                current = self._prepare_cloud(self.reader.point_cloud(storage_url), parameters)
                baseline_url = metadata.get("baseline_storage_url")
                if baseline_url:
                    baseline = self._prepare_cloud(self.reader.point_cloud(str(baseline_url)), parameters)
                    measurements = self._register_and_measure(baseline, current, parameters)
                    analysis = "Open3D Point-to-Plane ICP 与点云距离差分"
                else:
                    measurements = self._describe_cloud(current)
                    analysis = "Open3D 点云滤波与几何统计"
            lon = float(metadata.get("longitude", 116.102))
            lat = float(metadata.get("latitude", 39.102))
            confidence = 0.95 if measurements.get("quality_passed", True) else 0.55
            result = {
                "category": "点云形变" if baseline_url else "点云质量分析",
                "confidence": confidence,
                "geometry": {"type": "Point", "coordinates": [lon, lat]},
                "measurements": measurements,
                "attributes": {"analysis": analysis},
            }
            return SpatialOutcome(
                spec.runtime,
                [result],
                None,
                round((time.perf_counter() - started) * 1000, 2),
                self._model_info(spec),
            )
        except Exception as exc:
            return SpatialOutcome(
                "baseline",
                [],
                f"点云分析失败: {type(exc).__name__}: {exc}",
                round((time.perf_counter() - started) * 1000, 2),
                self._model_info(spec),
            )

    def status(self) -> dict[str, Any]:
        models = []
        for spec in self.registry.all():
            artifact = spec.artifact_path(self.settings.model_dir)
            is_geometry = spec.model_group == "pointcloud-analyzer"
            models.append(
                {
                    **self._model_info(spec),
                    "artifact_path": str(artifact) if artifact else None,
                    "artifact_installed": bool(o3d is not None) if is_geometry else bool(artifact and artifact.is_file()),
                    "loaded": self._session_key(spec) in self._sessions,
                    "fallback_available": is_geometry or cv2 is not None,
                }
            )
        return {
            "profile": self.registry.profile,
            "runtime": " / ".join(sorted({spec.runtime for spec in self.registry.all()})),
            "onnxruntime_available": ort is not None,
            "available_providers": ort.get_available_providers() if ort is not None else [],
            "execution_provider_ready": bool(ort is not None and self.settings.execution_provider in ort.get_available_providers()),
            "open3d_available": o3d is not None,
            "opencv_available": cv2 is not None,
            "execution_provider": self.settings.execution_provider,
            "intra_op_threads": self.settings.intra_op_threads,
            "inter_op_threads": self.settings.inter_op_threads,
            "max_concurrency": self.settings.max_concurrency,
            "max_loaded_models": self.settings.max_loaded_models,
            "max_point_count": self.settings.max_point_count,
            "fallback_mode": self.settings.fallback_mode,
            "models": models,
        }

    def unload(self, model_group: str | None = None) -> list[str]:
        with self._session_lock:
            if model_group:
                removed = [key for key in list(self._sessions) if key.startswith(model_group + ":")]
                for key in removed:
                    self._sessions.pop(key, None)
            else:
                removed = list(self._sessions.keys())
                self._sessions.clear()
        return removed

    def load(self, model_group: str, model_version: str | None = None) -> dict[str, Any]:
        spec = self.registry.get(model_group, model_version)
        if spec is None:
            raise ValueError("模型组或版本未在运行时注册表中配置")
        if spec.model_group == "pointcloud-analyzer":
            return {"model_group": spec.model_group, "model_version": spec.model_version, "loaded": True, "builtin": True}
        artifact = spec.artifact_path(self.settings.model_dir)
        if artifact is None or not artifact.is_file():
            raise FileNotFoundError(f"模型制品未安装: {artifact}")
        self.unload(spec.model_group)
        self._session(spec)
        return {"model_group": spec.model_group, "model_version": spec.model_version, "loaded": True}

    def _onnx_change(
        self, spec: SpatialModelSpec, baseline: np.ndarray, current: np.ndarray, parameters: dict[str, Any]
    ) -> list[dict[str, Any]]:
        session = self._session(spec)
        size = spec.input_size or 512
        first = self._image_tensor(baseline, size, spec)
        second = self._image_tensor(current, size, spec)
        inputs = session.get_inputs()
        if len(inputs) < 2:
            raise ValueError("BIT 变化检测模型必须包含两个图像输入")
        outputs = session.run(None, {inputs[0].name: first, inputs[1].name: second})
        array = max((np.asarray(item) for item in outputs), key=lambda item: item.size)
        if array.ndim == 4:
            array = array[0]
        if array.ndim == 3 and array.shape[0] > 1:
            mask = np.argmax(array, axis=0).astype(np.uint8)
            score_map = np.max(array, axis=0)
        elif array.ndim == 3:
            score_map = array[0]
            mask = (score_map >= float(parameters.get("change_threshold", 0.25))).astype(np.uint8)
        else:
            score_map = array
            mask = (array >= float(parameters.get("change_threshold", 0.25))).astype(np.uint8)
        return self._change_regions(mask, score_map, parameters)

    def _opencv_change(
        self, baseline: np.ndarray, current: np.ndarray, parameters: dict[str, Any]
    ) -> list[dict[str, Any]]:
        height = min(baseline.shape[0], current.shape[0])
        width = min(baseline.shape[1], current.shape[1])
        baseline = cv2.resize(baseline, (width, height))
        current = cv2.resize(current, (width, height))
        base_gray = cv2.cvtColor(baseline, cv2.COLOR_BGR2GRAY)
        current_gray = cv2.cvtColor(current, cv2.COLOR_BGR2GRAY)
        warp = np.eye(2, 3, dtype=np.float32)
        try:
            cv2.findTransformECC(base_gray, current_gray, warp, cv2.MOTION_AFFINE, (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 50, 1e-5))
            base_gray = cv2.warpAffine(base_gray, warp, (width, height), flags=cv2.INTER_LINEAR | cv2.WARP_INVERSE_MAP)
        except cv2.error:
            pass
        diff = cv2.absdiff(base_gray, current_gray).astype(np.float32) / 255.0
        threshold = float(parameters.get("change_threshold", 0.25))
        mask = (diff >= threshold).astype(np.uint8)
        kernel_size = int(parameters.get("morphology_kernel", 3))
        if kernel_size > 1:
            kernel = np.ones((kernel_size, kernel_size), dtype=np.uint8)
            mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
            mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        return self._change_regions(mask, diff, parameters)

    def _change_regions(
        self, mask: np.ndarray, score_map: np.ndarray, parameters: dict[str, Any]
    ) -> list[dict[str, Any]]:
        contours, _ = cv2.findContours((mask * 255).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        resolution = float(parameters.get("target_resolution_m", 0.1))
        minimum_pixels = float(parameters.get("minimum_area_m2", 2.0)) / max(1e-9, resolution * resolution)
        results: list[dict[str, Any]] = []
        for contour in sorted(contours, key=cv2.contourArea, reverse=True):
            area = cv2.contourArea(contour)
            if area < minimum_pixels:
                continue
            polygon = cv2.approxPolyDP(contour, 2.0, True)
            points = [[round(float(p[0][0]) / mask.shape[1], 6), round(float(p[0][1]) / mask.shape[0], 6)] for p in polygon]
            if len(points) < 3:
                continue
            points.append(points[0])
            region_mask = np.zeros(mask.shape, dtype=np.uint8)
            cv2.drawContours(region_mask, [contour], -1, 1, -1)
            confidence = float(np.mean(score_map[region_mask == 1])) if np.any(region_mask) else 0.0
            results.append(
                {
                    "category": "变化区域",
                    "confidence": round(min(1.0, max(0.0, confidence)), 4),
                    "geometry": {"type": "Polygon", "coordinates": [points], "coordinate_space": "normalized"},
                    "measurements": {"change_area_m2": round(area * resolution * resolution, 3)},
                    "attributes": {"pixel_area": int(area)},
                }
            )
        return results

    def _prepare_cloud(self, cloud: Any, parameters: dict[str, Any]) -> Any:
        points = np.asarray(cloud.points)
        if len(points) > self.settings.max_point_count:
            ratio = self.settings.max_point_count / len(points)
            cloud = cloud.random_down_sample(ratio)
        voxel_size = float(parameters.get("voxel_size_m", 0.08))
        if voxel_size > 0:
            cloud = cloud.voxel_down_sample(voxel_size)
        neighbors = min(int(parameters.get("outlier_neighbors", 20)), max(4, len(cloud.points) - 1))
        if len(cloud.points) > neighbors:
            cloud, _ = cloud.remove_statistical_outlier(
                nb_neighbors=neighbors,
                std_ratio=float(parameters.get("outlier_std_ratio", 2.0)),
            )
        return cloud

    def _register_and_measure(self, baseline: Any, current: Any, parameters: dict[str, Any]) -> dict[str, Any]:
        if min(len(baseline.points), len(current.points)) < 20:
            raise ValueError("点云有效点数不足，无法配准")
        radius = max(float(parameters.get("voxel_size_m", 0.08)) * 3, 0.05)
        baseline.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=radius, max_nn=30))
        current.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=radius, max_nn=30))
        threshold = float(parameters.get("icp_max_distance_m", 0.5))
        criteria = o3d.pipelines.registration.ICPConvergenceCriteria(
            max_iteration=int(parameters.get("icp_iterations", 80))
        )
        registration = o3d.pipelines.registration.registration_icp(
            current,
            baseline,
            threshold,
            np.eye(4),
            o3d.pipelines.registration.TransformationEstimationPointToPlane(),
            criteria,
        )
        aligned = current.transform(registration.transformation.copy())
        distances = np.asarray(aligned.compute_point_cloud_distance(baseline))
        if distances.size == 0:
            raise ValueError("配准后没有可比较点")
        change_threshold = float(parameters.get("registration_error_limit_m", 0.1))
        changed = distances > change_threshold
        voxel = float(parameters.get("voxel_size_m", 0.08))
        volume_proxy = float(np.sum(np.maximum(0, distances[changed] - change_threshold)) * voxel * voxel)
        quality_passed = registration.inlier_rmse <= change_threshold and registration.fitness >= 0.25
        return {
            "point_count": int(distances.size),
            "registration_error_m": round(float(registration.inlier_rmse), 6),
            "registration_fitness": round(float(registration.fitness), 6),
            "mean_distance_m": round(float(np.mean(distances)), 6),
            "p95_distance_m": round(float(np.percentile(distances, 95)), 6),
            "max_height_diff_m": round(float(np.max(distances)), 6),
            "changed_point_count": int(np.sum(changed)),
            "changed_point_ratio": round(float(np.mean(changed)), 6),
            "volume_change_m3": round(volume_proxy, 6),
            "quality_passed": bool(quality_passed),
        }

    def _describe_cloud(self, cloud: Any) -> dict[str, Any]:
        points = np.asarray(cloud.points)
        minimum = points.min(axis=0)
        maximum = points.max(axis=0)
        extent = maximum - minimum
        return {
            "point_count": int(len(points)),
            "extent_x_m": round(float(extent[0]), 4),
            "extent_y_m": round(float(extent[1]), 4),
            "extent_z_m": round(float(extent[2]), 4),
            "min_elevation_m": round(float(minimum[2]), 4),
            "max_elevation_m": round(float(maximum[2]), 4),
            "quality_passed": len(points) >= 1000,
        }

    def _session(self, spec: SpatialModelSpec) -> Any:
        if ort is None:
            raise RuntimeError("onnxruntime 未安装")
        with self._session_lock:
            key = self._session_key(spec)
            existing = self._sessions.pop(key, None)
            if existing is not None:
                self._sessions[key] = existing
                return existing
            options = ort.SessionOptions()
            options.intra_op_num_threads = self.settings.intra_op_threads
            options.inter_op_num_threads = self.settings.inter_op_threads
            options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
            options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
            session = ort.InferenceSession(
                str(spec.artifact_path(self.settings.model_dir)),
                sess_options=options,
                providers=[self.settings.execution_provider],
            )
            while len(self._sessions) >= self.settings.max_loaded_models:
                self._sessions.popitem(last=False)
            self._sessions[key] = session
            return session

    def _session_key(self, spec: SpatialModelSpec) -> str:
        return f"{spec.model_group}:{spec.model_version}"

    def _image_tensor(self, image: np.ndarray, size: int, spec: SpatialModelSpec) -> np.ndarray:
        image = cv2.cvtColor(cv2.resize(image, (size, size)), cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
        image = (image - np.asarray(spec.mean, dtype=np.float32)) / np.asarray(spec.std, dtype=np.float32)
        return np.transpose(image, (2, 0, 1))[None, ...]

    def _model_info(self, spec: SpatialModelSpec) -> dict[str, Any]:
        return {
            "model_group": spec.model_group,
            "model_version": spec.model_version,
            "display_name": spec.display_name,
            "family": spec.family,
            "runtime": spec.runtime,
            "precision": spec.precision,
            "input_size": spec.input_size,
        }

    def _unavailable(self, model_group: str, reason: str, model_version: str | None = None) -> SpatialOutcome:
        spec = self.registry.get(model_group, model_version)
        model = self._model_info(spec) if spec else {"model_group": model_group}
        return SpatialOutcome("baseline", [], reason, 0.0, model)
