"""Model Manager (Component F16).

Backends:
  - MockInferenceEngine    — NumPy stubs, works everywhere (dev/CI)
  - TRTInferenceEngine     — TensorRT FP16 engine loader (Jetson only, VO-03)

ModelManager.get_inference_engine() auto-selects TRT on Jetson when a .engine
file is available, otherwise falls back to Mock.
"""

import logging
import os
from abc import ABC, abstractmethod
from typing import Any

import numpy as np

from gps_denied.schemas.model import InferenceEngine

logger = logging.getLogger(__name__)


def _is_jetson() -> bool:
    """Return True when running on an NVIDIA Jetson device."""
    try:
        with open("/proc/device-tree/compatible", "rb") as f:
            return b"nvidia,tegra" in f.read()
    except OSError:
        pass
    # Secondary check: tegra chip_id
    return os.path.exists("/sys/bus/platform/drivers/tegra-se-nvhost")


class IModelManager(ABC):
    @abstractmethod
    def load_model(self, model_name: str, model_format: str) -> bool:
        pass

    @abstractmethod
    def get_inference_engine(self, model_name: str) -> InferenceEngine:
        pass

    @abstractmethod
    def optimize_to_tensorrt(self, model_name: str, onnx_path: str) -> str:
        pass

    @abstractmethod
    def fallback_to_onnx(self, model_name: str) -> bool:
        pass

    @abstractmethod
    def warmup_model(self, model_name: str) -> bool:
        pass


class MockInferenceEngine(InferenceEngine):
    """A mock implementation of Inference Engine for rapid testing."""
    def infer(self, input_data: Any) -> Any:
        if self.model_name == "SuperPoint":
            # Mock extracting 500 features
            n_features = 500
            # Assuming input_data is an image of shape (H, W, 3)
            h, w = input_data.shape[:2] if hasattr(input_data, "shape") else (480, 640)

            keypoints = np.random.rand(n_features, 2) * [w, h]
            descriptors = np.random.rand(n_features, 256)
            scores = np.random.rand(n_features)

            return {
                "keypoints": keypoints,
                "descriptors": descriptors,
                "scores": scores
            }

        elif self.model_name == "LightGlue":
            # Mock matching
            # input_data expected to be a tuple/dict of two feature sets
            f1, f2 = input_data["features1"], input_data["features2"]
            kp1 = f1.keypoints
            kp2 = f2.keypoints

            # Create ~100 random matches
            n_matches = min(100, len(kp1), len(kp2))

            indices1 = np.random.choice(len(kp1), n_matches, replace=False)
            indices2 = np.random.choice(len(kp2), n_matches, replace=False)

            matches = np.stack([indices1, indices2], axis=1)
            scores = np.random.rand(n_matches)

            return {
                "matches": matches,
                "scores": scores,
                "keypoints1": kp1[indices1],
                "keypoints2": kp2[indices2]
            }

        elif self.model_name == "DINOv2":
            # Mock generating 4096-dim VLAD descriptor
            dim = 4096
            desc = np.random.rand(dim)
            # L2 normalize
            return desc / np.linalg.norm(desc)

        elif self.model_name in ("LiteSAM", "XFeat"):
            # Mock LiteSAM / XFeat matching between UAV and satellite image.
            # Returns homography, inlier_count, total_correspondences, confidence.
            homography = np.eye(3, dtype=np.float64)
            homography[0, 2] = np.random.uniform(-50, 50)
            homography[1, 2] = np.random.uniform(-50, 50)

            # 80% chance to produce a good match
            matched = np.random.rand() > 0.2
            total = np.random.randint(80, 200)
            inliers = np.random.randint(40, total) if matched else np.random.randint(0, 15)

            return {
                "homography": homography,
                "inlier_count": inliers,
                "total_correspondences": total,
                "confidence": inliers / max(total, 1),
                "reprojection_error": np.random.uniform(0.3, 1.5) if matched else 5.0,
            }

        raise ValueError(f"Unknown mock model: {self.model_name}")


class TRTInferenceEngine(InferenceEngine):
    """TensorRT FP16 inference engine — Jetson only (VO-03).

    Loads a pre-built .engine file produced by trtexec --fp16.
    Falls back to MockInferenceEngine if TensorRT is unavailable.
    """

    def __init__(self, model_name: str, engine_path: str):
        super().__init__(model_name, "trt")
        self._engine_path = engine_path
        self._runtime = None
        self._engine = None
        self._context = None
        self._mock_fallback: MockInferenceEngine | None = None
        self._load()

    def _load(self):
        try:
            import pycuda.autoinit  # type: ignore  # noqa: F401
            import pycuda.driver  # type: ignore  # noqa: F401
            import tensorrt as trt  # type: ignore

            trt_logger = trt.Logger(trt.Logger.WARNING)
            self._runtime = trt.Runtime(trt_logger)
            with open(self._engine_path, "rb") as f:
                self._engine = self._runtime.deserialize_cuda_engine(f.read())
            self._context = self._engine.create_execution_context()
            logger.info("TRTInferenceEngine: loaded %s from %s", self.model_name, self._engine_path)
        except (ImportError, FileNotFoundError, Exception) as exc:
            logger.info(
                "TRTInferenceEngine: cannot load %s (%s) — using Mock", self.model_name, exc
            )
            self._mock_fallback = MockInferenceEngine(self.model_name, "mock")

    def infer(self, input_data: Any) -> Any:
        if self._mock_fallback is not None:
            return self._mock_fallback.infer(input_data)
        # Real TRT inference — placeholder for host↔device transfer logic
        raise NotImplementedError(
            "Real TRT inference not yet wired — provide a model-specific subclass"
        )


class ModelManager(IModelManager):
    """Manages ML models lifecycle and provisioning.

    On Jetson (cuDA/TRT available) and when a matching .engine file exists under
    `engine_dir`, loads TRTInferenceEngine.  Otherwise uses MockInferenceEngine.
    """

    # Map model name → expected .engine filename
    _TRT_ENGINE_FILES: dict[str, str] = {
        "SuperPoint": "superpoint.engine",
        "LightGlue": "lightglue.engine",
        "XFeat": "xfeat.engine",
        "DINOv2": "dinov2.engine",
        "LiteSAM": "litesam.engine",
    }

    def __init__(self, engine_dir: str = "/opt/engines"):
        self._loaded_models: dict[str, InferenceEngine] = {}
        self._engine_dir = engine_dir
        self._on_jetson = _is_jetson()

    def _engine_path(self, model_name: str) -> str | None:
        """Return full path to .engine file if it exists, else None."""
        filename = self._TRT_ENGINE_FILES.get(model_name)
        if filename is None:
            return None
        path = os.path.join(self._engine_dir, filename)
        return path if os.path.isfile(path) else None

    def load_model(self, model_name: str, model_format: str) -> bool:
        """Load a model. Uses TRT on Jetson when engine file exists, Mock otherwise."""
        logger.info("Loading %s (format=%s, jetson=%s)", model_name, model_format, self._on_jetson)

        engine_path = self._engine_path(model_name) if self._on_jetson else None
        if engine_path:
            engine: InferenceEngine = TRTInferenceEngine(model_name, engine_path)
        else:
            engine = MockInferenceEngine(model_name, model_format)

        self._loaded_models[model_name] = engine
        self.warmup_model(model_name)
        return True

    def get_inference_engine(self, model_name: str) -> InferenceEngine:
        """Gets an inference engine, auto-loading if needed."""
        if model_name not in self._loaded_models:
            self.load_model(model_name, "trt" if self._on_jetson else "mock")
        return self._loaded_models[model_name]

    def optimize_to_tensorrt(self, model_name: str, onnx_path: str) -> str:
        """Placeholder for TensorRT optimization."""
        return f"{onnx_path}.trt"

    def fallback_to_onnx(self, model_name: str) -> bool:
        """Placeholder for fallback logic."""
        logger.warning(f"Falling back to ONNX for {model_name}")
        return True

    def warmup_model(self, model_name: str) -> bool:
        """Warms up a loaded model."""
        logger.info(f"Warming up {model_name}")
        return True
