from __future__ import annotations

import base64
import json
import os
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

from .schemas import Resource

try:
    import onnxruntime as ort
except ImportError:  # pragma: no cover - exercised by the lightweight source test environment
    ort = None

try:
    from minio import Minio
except ImportError:  # pragma: no cover - exercised by the lightweight source test environment
    Minio = None


@dataclass(frozen=True)
class CpuRuntimeSettings:
    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
    fallback_mode: str

    @classmethod
    def from_env(cls) -> "CpuRuntimeSettings":
        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", "64"))) * 1024 * 1024,
            fallback_mode=os.getenv("RAIL_FALLBACK_MODE", "baseline"),
        )


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

    @classmethod
    def from_dict(cls, value: dict[str, Any]) -> "ModelSpec":
        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"),
            secondary_artifact=value.get("secondary_artifact"),
            labels=value.get("labels"),
            parser=str(value.get("parser", "paddle-detection")),
            input_size=int(value["input_size"]) if value.get("input_size") else None,
            batch_size=int(value.get("batch_size", 1)),
            precision=str(value.get("precision", "FP32")),
            mean=tuple(float(item) for item in value.get("mean", [0.0, 0.0, 0.0])),
            std=tuple(float(item) for item in value.get("std", [1.0, 1.0, 1.0])),
            scale=float(value.get("scale", 1 / 255.0)),
            color_order=str(value.get("color_order", "RGB")),
            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

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

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


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


class ModelRegistry:
    def __init__(self, settings: CpuRuntimeSettings):
        self.settings = settings
        payload = self._load_payload()
        self.profile = str(payload.get("profile", settings.profile))
        self._models: dict[str, list[ModelSpec]] = {}
        for spec in (ModelSpec.from_dict(item) for item in payload.get("models", [])):
            if spec.model_group in {"vision-detector", "vision-segmenter", "thermal-analyzer"}:
                self._models.setdefault(spec.model_group, []).append(spec)

    def get(self, model_group: str, model_version: str | None = None) -> ModelSpec | None:
        specs = self._models.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[ModelSpec]:
        return [spec for specs in self._models.values() for spec in specs]

    def _load_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": "vision-detector",
                    "model_version": "cpu-v1.0.0",
                    "display_name": "PP-YOLOE-SOD-s CPU",
                    "family": "PP-YOLOE-SOD",
                    "artifact": "vision-detector/ppyoloe-sod-s/model.onnx",
                    "labels": "vision-detector/ppyoloe-sod-s/labels.txt",
                    "parser": "paddle-detection",
                    "input_size": 640,
                },
                {
                    "model_group": "vision-segmenter",
                    "model_version": "cpu-v1.0.0",
                    "display_name": "PP-LiteSeg-STDC1 CPU",
                    "family": "PP-LiteSeg",
                    "artifact": "vision-segmenter/pp-liteseg-stdc1/model.onnx",
                    "labels": "vision-segmenter/pp-liteseg-stdc1/labels.txt",
                    "parser": "semantic-segmentation",
                    "input_size": 512,
                    "mean": [0.5, 0.5, 0.5],
                    "std": [0.5, 0.5, 0.5],
                },
                {
                    "model_group": "thermal-analyzer",
                    "model_version": "cpu-v1.0.0",
                    "display_name": "PP-YOLOE-s + PatchCore CPU",
                    "family": "PP-YOLOE/PatchCore",
                    "artifact": "thermal-analyzer/ppyoloe-s/model.onnx",
                    "labels": "thermal-analyzer/ppyoloe-s/labels.txt",
                    "parser": "thermal-detection",
                    "input_size": 512,
                },
            ],
        }


class ResourceLoader:
    def __init__(self, settings: CpuRuntimeSettings):
        self.settings = settings
        self._minio: Any = None

    def image(self, resource: Resource) -> np.ndarray:
        raw = self.bytes(resource)
        image = cv2.imdecode(np.frombuffer(raw, dtype=np.uint8), cv2.IMREAD_UNCHANGED)
        if image is None:
            raise ValueError("输入文件不是可解码的图像")
        if image.ndim == 2:
            image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
        if image.shape[2] == 4:
            image = cv2.cvtColor(image, cv2.COLOR_BGRA2BGR)
        return image

    def bytes(self, resource: Resource) -> bytes:
        encoded = resource.metadata.get("image_base64")
        if encoded:
            payload = base64.b64decode(str(encoded), validate=True)
            return self._bounded(payload)
        location = resource.storage_url or resource.url
        if not location:
            raise FileNotFoundError("资源没有可读取的 storage_url 或 url")
        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=15) as response:  # noqa: S310 - endpoint is controlled by deployment config
                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 _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 LazyOnnxRuntime:
    def __init__(self, settings: CpuRuntimeSettings, registry: ModelRegistry):
        self.settings = settings
        self.registry = registry
        self._sessions: OrderedDict[str, Any] = OrderedDict()
        self._lock = threading.RLock()
        self._semaphore = threading.BoundedSemaphore(settings.max_concurrency)

    def predict(
        self,
        model_group: str,
        image: np.ndarray,
        parameters: dict[str, Any],
        model_version: str | None = None,
    ) -> RuntimeOutcome:
        spec = self.registry.get(model_group, model_version)
        if spec is None:
            return self._unavailable(model_group, "模型组或版本未在运行时注册表中配置", model_version)
        artifact = spec.artifact_path(self.settings.model_dir)
        if artifact is None or not artifact.is_file():
            return self._unavailable(model_group, f"模型制品未安装: {artifact}", model_version)
        if ort is None:
            return self._unavailable(model_group, "onnxruntime 未安装", model_version)

        started = time.perf_counter()
        try:
            with self._semaphore:
                session = self._session(spec)
                tensor, transform = self._preprocess(image, spec, session)
                feeds = self._feeds(session, tensor, transform)
                outputs = session.run(None, feeds)
                results = self._parse(outputs, spec, transform, parameters)
            return RuntimeOutcome(
                execution_mode=spec.runtime,
                results=results,
                reason=None,
                latency_ms=round((time.perf_counter() - started) * 1000, 2),
                model=self._model_info(spec, True),
            )
        except Exception as exc:  # model contracts differ; expose the precise reason and use governed fallback
            return RuntimeOutcome(
                execution_mode="baseline",
                results=[],
                reason=f"CPU 模型推理失败: {type(exc).__name__}: {exc}",
                latency_ms=round((time.perf_counter() - started) * 1000, 2),
                model=self._model_info(spec, False),
            )

    def status(self) -> dict[str, Any]:
        loaded = set(self._sessions.keys())
        models = []
        for spec in self.registry.all():
            artifact = spec.artifact_path(self.settings.model_dir)
            secondary = spec.secondary_artifact_path(self.settings.model_dir)
            labels = spec.labels_path(self.settings.model_dir)
            artifacts_ready = bool(artifact and artifact.is_file()) and bool(secondary is None or secondary.is_file())
            models.append(
                {
                    **self._model_info(spec, artifacts_ready),
                    "artifact_path": str(artifact) if artifact else None,
                    "secondary_artifact_path": str(secondary) if secondary else None,
                    "secondary_artifact_installed": secondary.is_file() if secondary else None,
                    "labels_installed": labels.is_file() if labels else False,
                    "loaded": self._session_key(spec) in loaded,
                }
            )
        return {
            "profile": self.registry.profile,
            "runtime": " / ".join(sorted({spec.runtime for spec in self.registry.all()})),
            "runtime_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()),
            "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,
            "fallback_mode": self.settings.fallback_mode,
            "models": models,
        }

    def unload(self, model_group: str | None = None) -> list[str]:
        with self._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("模型组或版本未在运行时注册表中配置")
        artifact = spec.artifact_path(self.settings.model_dir)
        secondary = spec.secondary_artifact_path(self.settings.model_dir)
        if artifact is None or not artifact.is_file():
            raise FileNotFoundError(f"模型制品未安装: {artifact}")
        if secondary is not None and not secondary.is_file():
            raise FileNotFoundError(f"辅助模型制品未安装: {secondary}")
        self.unload(spec.model_group)
        self._session(spec)
        return {"model_group": spec.model_group, "model_version": spec.model_version, "loaded": True}

    def _session(self, spec: ModelSpec) -> Any:
        with self._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: ModelSpec) -> str:
        return f"{spec.model_group}:{spec.model_version}"

    def _preprocess(self, image: np.ndarray, spec: ModelSpec, session: Any) -> tuple[np.ndarray, dict[str, Any]]:
        original_h, original_w = image.shape[:2]
        input_meta = session.get_inputs()[0]
        shape = input_meta.shape
        shape_h = shape[-2] if len(shape) >= 4 and isinstance(shape[-2], int) else None
        shape_w = shape[-1] if len(shape) >= 4 and isinstance(shape[-1], int) else None
        target_h = int(shape_h or spec.input_size or 640)
        target_w = int(shape_w or spec.input_size or 640)
        resized = cv2.resize(image, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
        if spec.color_order.upper() == "RGB":
            resized = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
        tensor = resized.astype(np.float32) * spec.scale
        tensor = (tensor - np.asarray(spec.mean, dtype=np.float32)) / np.asarray(spec.std, dtype=np.float32)
        tensor = np.transpose(tensor, (2, 0, 1))[None, ...]
        return tensor, {
            "original_h": original_h,
            "original_w": original_w,
            "input_h": target_h,
            "input_w": target_w,
            "scale_y": target_h / original_h,
            "scale_x": target_w / original_w,
        }

    def _feeds(self, session: Any, tensor: np.ndarray, transform: dict[str, Any]) -> dict[str, np.ndarray]:
        feeds: dict[str, np.ndarray] = {}
        for item in session.get_inputs():
            name = item.name
            lowered = name.lower()
            if "scale" in lowered:
                feeds[name] = np.asarray([[transform["scale_y"], transform["scale_x"]]], dtype=np.float32)
            elif "shape" in lowered or "size" in lowered:
                feeds[name] = np.asarray([[transform["input_h"], transform["input_w"]]], dtype=np.float32)
            else:
                feeds[name] = tensor
        return feeds

    def _parse(
        self,
        outputs: list[np.ndarray],
        spec: ModelSpec,
        transform: dict[str, Any],
        parameters: dict[str, Any],
    ) -> list[dict[str, Any]]:
        if spec.parser == "semantic-segmentation":
            return self._parse_segmentation(outputs, spec, parameters)
        return self._parse_detection(outputs, spec, transform, parameters)

    def _parse_detection(
        self,
        outputs: list[np.ndarray],
        spec: ModelSpec,
        transform: dict[str, Any],
        parameters: dict[str, Any],
    ) -> list[dict[str, Any]]:
        threshold = float(parameters.get("confidence_threshold", 0.55))
        labels = self._labels(spec)
        candidates: np.ndarray | None = None
        for output in outputs:
            array = np.asarray(output)
            if array.ndim == 3 and array.shape[0] == 1:
                array = array[0]
            if array.ndim == 2 and array.shape[-1] >= 6:
                candidates = array
                if array.shape[-1] == 6:
                    break
        if candidates is None:
            raise ValueError("模型输出中没有可识别的检测矩阵")

        results: list[dict[str, Any]] = []
        max_detections = int(parameters.get("max_detections", 100))
        for row in candidates:
            if len(row) == 6:
                class_id, score, x1, y1, x2, y2 = row.tolist()
            else:
                objectness = float(row[4])
                class_scores = row[5:]
                class_id = int(np.argmax(class_scores))
                score = objectness * float(class_scores[class_id])
                cx, cy, width, height = (float(item) for item in row[:4])
                x1, y1, x2, y2 = cx - width / 2, cy - height / 2, cx + width / 2, cy + height / 2
            if float(score) < threshold:
                continue
            coordinates = [
                self._clip(float(x1) / transform["input_w"]),
                self._clip(float(y1) / transform["input_h"]),
                self._clip(float(x2) / transform["input_w"]),
                self._clip(float(y2) / transform["input_h"]),
            ]
            category = labels[int(class_id)] if 0 <= int(class_id) < len(labels) else f"class-{int(class_id)}"
            results.append(
                {
                    "category": category,
                    "confidence": round(float(score), 4),
                    "geometry": {"type": "BBox", "coordinates": coordinates, "coordinate_space": "normalized"},
                    "measurements": {},
                }
            )
            if len(results) >= max_detections:
                break
        return results

    def _parse_segmentation(
        self, outputs: list[np.ndarray], spec: ModelSpec, parameters: dict[str, Any]
    ) -> list[dict[str, Any]]:
        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)
        elif array.ndim == 3:
            threshold = float(parameters.get("mask_threshold", 0.5))
            mask = (array[0] >= threshold).astype(np.uint8)
        elif array.ndim == 2:
            mask = array.astype(np.uint8)
        else:
            raise ValueError(f"不支持的分割输出形状: {array.shape}")
        labels = self._labels(spec)
        minimum_area = int(parameters.get("minimum_area", 64))
        results: list[dict[str, Any]] = []
        for class_id in (int(item) for item in np.unique(mask) if int(item) != 0):
            binary = np.where(mask == class_id, 255, 0).astype(np.uint8)
            contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            for contour in sorted(contours, key=cv2.contourArea, reverse=True):
                area = cv2.contourArea(contour)
                if area < minimum_area:
                    continue
                epsilon = float(parameters.get("contour_simplification", 0.8))
                simplified = cv2.approxPolyDP(contour, max(0.5, epsilon), True)
                points = [
                    [self._clip(float(point[0][0]) / mask.shape[1]), self._clip(float(point[0][1]) / mask.shape[0])]
                    for point in simplified
                ]
                if len(points) < 3:
                    continue
                points.append(points[0])
                category = labels[class_id] if class_id < len(labels) else f"class-{class_id}"
                results.append(
                    {
                        "category": category,
                        "confidence": 1.0,
                        "geometry": {"type": "Polygon", "coordinates": [points], "coordinate_space": "normalized"},
                        "measurements": {"area_ratio": round(area / mask.size, 6)},
                    }
                )
        return results

    def _labels(self, spec: ModelSpec) -> list[str]:
        path = spec.labels_path(self.settings.model_dir)
        if path and path.is_file():
            return [line.strip() for line in path.read_text(encoding="utf-8-sig").splitlines() if line.strip()]
        return []

    def _model_info(self, spec: ModelSpec, artifact_installed: bool) -> 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,
            "artifact_installed": artifact_installed,
        }

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

    @staticmethod
    def _clip(value: float) -> float:
        return round(min(1.0, max(0.0, value)), 6)


class CpuVisionRuntime:
    def __init__(self):
        self.settings = CpuRuntimeSettings.from_env()
        self.registry = ModelRegistry(self.settings)
        self.resources = ResourceLoader(self.settings)
        self.models = LazyOnnxRuntime(self.settings, self.registry)

    def infer(
        self,
        resource: Resource,
        model_group: str,
        parameters: dict[str, Any],
        model_version: str | None = None,
    ) -> RuntimeOutcome:
        spec = self.registry.get(model_group, model_version)
        if spec is None:
            return RuntimeOutcome("baseline", [], "模型组未配置", 0.0, {"model_group": model_group})
        artifact = spec.artifact_path(self.settings.model_dir)
        if artifact is None or not artifact.is_file():
            return self.models.predict(model_group, np.empty((1, 1, 3), dtype=np.uint8), parameters, model_version)
        started = time.perf_counter()
        try:
            image = self.resources.image(resource)
        except Exception as exc:
            return RuntimeOutcome(
                "baseline",
                [],
                f"无法读取真实推理输入: {type(exc).__name__}: {exc}",
                round((time.perf_counter() - started) * 1000, 2),
                self.models._model_info(spec, True),
            )
        return self.models.predict(model_group, image, parameters, model_version)

    def status(self) -> dict[str, Any]:
        return self.models.status()

    def load(self, model_group: str, model_version: str | None = None) -> dict[str, Any]:
        return self.models.load(model_group, model_version)
