"""
Edge Inference Runtime — wraps VL-JEPA for edge deployment on GPU boxes.

Ties together the full pipeline:
    RTSP cameras → KeyframeSampler → Preprocessing → Encoder → SelectiveDecoder → MoE Decoder

Key responsibilities:
- Frame preprocessing: 720p RTSP → resize 384x384 → normalize → tensor
- Selective decode integration: only run MoE decoder when semantic shift detected
- Batch inference across multiple cameras
- Performance metrics: fps, latency, decode ratio
- ONNX export helper for future TensorRT optimization

The runtime is designed for Jetson / edge GPU boxes managing 4-64 cameras,
where compute budget is tight and selective decoding is essential.
"""

import cv2
import time
import logging
import threading
from collections import defaultdict, deque
from dataclasses import dataclass, field
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from model.vlm import VLJEPAModel
from model.selective_decode import SelectiveDecoder
from edge.ingest import CameraManager, RTSPCamera
from edge.sampler import MultiCameraSampler, KeyframeSampler

logger = logging.getLogger(__name__)

# ImageNet normalization constants (used by most vision models)
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]


# ---------------------------------------------------------------------------
# Preprocessing
# ---------------------------------------------------------------------------

class FramePreprocessor:
    """
    Converts raw BGR camera frames to model-ready tensors.

    Pipeline: BGR 720p → RGB → resize 384x384 → float32 [0,1] → normalize → CHW tensor

    Args:
        target_size: (H, W) input size expected by the ViT encoder
        mean: Per-channel mean for normalization
        std: Per-channel std for normalization
        device: Target torch device
    """

    def __init__(
        self,
        target_size: tuple[int, int] = (384, 384),
        mean: list[float] = None,
        std: list[float] = None,
        device: str = "cuda",
    ):
        self.target_size = target_size
        self.mean = np.array(mean or IMAGENET_MEAN, dtype=np.float32).reshape(1, 1, 3)
        self.std = np.array(std or IMAGENET_STD, dtype=np.float32).reshape(1, 1, 3)
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")

    def preprocess(self, frame: np.ndarray) -> torch.Tensor:
        """
        Single frame preprocessing.

        Args:
            frame: BGR uint8 image from OpenCV (any resolution)
        Returns:
            [1, 3, 384, 384] float32 tensor on target device
        """
        # BGR → RGB
        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        # Resize to model input size
        resized = cv2.resize(rgb, self.target_size, interpolation=cv2.INTER_LINEAR)
        # float32 [0, 1]
        normalized = resized.astype(np.float32) / 255.0
        # ImageNet normalization
        normalized = (normalized - self.mean) / self.std
        # HWC → CHW → BCHW
        tensor = torch.from_numpy(normalized.transpose(2, 0, 1)).unsqueeze(0)
        return tensor.to(self.device)

    def preprocess_batch(self, frames: list[np.ndarray]) -> torch.Tensor:
        """
        Batch preprocessing for multiple frames.

        Args:
            frames: List of BGR uint8 images
        Returns:
            [B, 3, 384, 384] float32 tensor
        """
        if len(frames) == 0:
            return torch.empty(0, 3, *self.target_size, device=self.device)

        tensors = []
        for frame in frames:
            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            resized = cv2.resize(rgb, self.target_size, interpolation=cv2.INTER_LINEAR)
            normalized = resized.astype(np.float32) / 255.0
            normalized = (normalized - self.mean) / self.std
            tensors.append(normalized.transpose(2, 0, 1))

        batch = np.stack(tensors, axis=0)
        return torch.from_numpy(batch).to(self.device)


# ---------------------------------------------------------------------------
# Performance metrics
# ---------------------------------------------------------------------------

@dataclass
class InferenceMetrics:
    """Tracks runtime performance statistics."""
    # Latency tracking (sliding window)
    _encode_latencies: deque = field(default_factory=lambda: deque(maxlen=100))
    _decode_latencies: deque = field(default_factory=lambda: deque(maxlen=100))
    _total_latencies: deque = field(default_factory=lambda: deque(maxlen=100))
    _frame_times: deque = field(default_factory=lambda: deque(maxlen=100))

    # Counters
    frames_processed: int = 0
    decodes_triggered: int = 0
    decodes_skipped: int = 0

    def record_encode(self, latency_sec: float) -> None:
        self._encode_latencies.append(latency_sec * 1000)

    def record_decode(self, latency_sec: float) -> None:
        self._decode_latencies.append(latency_sec * 1000)

    def record_total(self, latency_sec: float) -> None:
        self._total_latencies.append(latency_sec * 1000)
        now = time.monotonic()
        self._frame_times.append(now)
        self.frames_processed += 1

    @property
    def fps(self) -> float:
        """Effective processing throughput."""
        if len(self._frame_times) < 2:
            return 0.0
        elapsed = self._frame_times[-1] - self._frame_times[0]
        if elapsed <= 0:
            return 0.0
        return (len(self._frame_times) - 1) / elapsed

    @property
    def encode_latency_ms(self) -> float:
        if not self._encode_latencies:
            return 0.0
        return sum(self._encode_latencies) / len(self._encode_latencies)

    @property
    def decode_latency_ms(self) -> float:
        if not self._decode_latencies:
            return 0.0
        return sum(self._decode_latencies) / len(self._decode_latencies)

    @property
    def total_latency_ms(self) -> float:
        if not self._total_latencies:
            return 0.0
        return sum(self._total_latencies) / len(self._total_latencies)

    @property
    def decode_ratio(self) -> float:
        total = self.decodes_triggered + self.decodes_skipped
        if total == 0:
            return 0.0
        return self.decodes_triggered / total

    def to_dict(self) -> dict:
        return {
            "fps": round(self.fps, 2),
            "encode_latency_ms": round(self.encode_latency_ms, 2),
            "decode_latency_ms": round(self.decode_latency_ms, 2),
            "total_latency_ms": round(self.total_latency_ms, 2),
            "frames_processed": self.frames_processed,
            "decodes_triggered": self.decodes_triggered,
            "decodes_skipped": self.decodes_skipped,
            "decode_ratio": round(self.decode_ratio, 4),
        }


# ---------------------------------------------------------------------------
# Edge Inference Server
# ---------------------------------------------------------------------------

class EdgeInferenceServer:
    """
    Wraps the VL-JEPA model for edge deployment with multi-camera support.

    Orchestrates the full pipeline:
        CameraManager → MultiCameraSampler → FramePreprocessor
        → VLJEPAModel.get_embedding → SelectiveDecoder.should_decode
        → VLJEPAModel.decoder (only on semantic shift) → text output

    The server runs a processing loop in a background thread, pulling
    keyframes from all cameras, running inference, and storing results.

    Args:
        model: Loaded VLJEPAModel instance
        device: Torch device string ("cuda", "cuda:0", "cpu")
        selective_threshold: Cosine similarity threshold for selective decoding
        min_decode_interval: Minimum seconds between decodes per camera
        max_new_tokens: Max tokens for text generation
        temperature: Sampling temperature for generation
    """

    def __init__(
        self,
        model: VLJEPAModel,
        device: str = "cuda",
        selective_threshold: float = 0.95,
        min_decode_interval: float = 1.0,
        max_new_tokens: int = 128,
        temperature: float = 0.8,
    ):
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
        self.model = model.to(self.device).eval()
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature

        # Components
        self.camera_manager = CameraManager()
        self.sampler = MultiCameraSampler()
        self.preprocessor = FramePreprocessor(device=str(self.device))
        self.selective_decoder = SelectiveDecoder(
            similarity_threshold=selective_threshold,
            min_decode_interval=min_decode_interval,
            embed_dim=model.selective_decoder.embed_dim,
        )

        # Results storage: camera_id → latest generation result
        self._results: dict[str, dict] = {}
        self._results_lock = threading.Lock()

        # Metrics
        self.metrics = InferenceMetrics()

        # Processing loop control
        self._thread: Optional[threading.Thread] = None
        self._stop_event = threading.Event()
        self._processing_interval = 0.05  # 50ms between processing cycles

    # ------------------------------------------------------------------
    # Camera management (delegates to CameraManager)
    # ------------------------------------------------------------------

    def add_camera(
        self,
        camera_id: str,
        rtsp_url: str,
        target_fps: float = 5.0,
        **kwargs,
    ) -> None:
        """Register and start an RTSP camera."""
        self.camera_manager.add_camera(camera_id, rtsp_url, target_fps=target_fps, **kwargs)
        self.camera_manager.start_camera(camera_id)
        logger.info(f"Camera '{camera_id}' added and started")

    def remove_camera(self, camera_id: str) -> None:
        """Stop and remove a camera."""
        self.camera_manager.remove_camera(camera_id)
        self.sampler.remove_camera(camera_id)
        self.selective_decoder.reset_camera(camera_id)
        with self._results_lock:
            self._results.pop(camera_id, None)

    # ------------------------------------------------------------------
    # Single-frame inference
    # ------------------------------------------------------------------

    @torch.no_grad()
    def process_frame(
        self,
        camera_id: str,
        frame: np.ndarray,
        timestamp: float,
        force_decode: bool = False,
    ) -> Optional[dict]:
        """
        Process a single frame through the full pipeline.

        Steps:
            1. Preprocess frame → tensor
            2. Run encoder → get embedding
            3. Check selective decoder → should we decode?
            4. If yes (or force_decode), run MoE decoder → text

        Args:
            camera_id: Camera identifier
            frame: BGR image
            timestamp: Wall-clock time
            force_decode: Bypass selective decoder (e.g., for user queries)
        Returns:
            Dict with results if decode was triggered, else None.
            Keys: camera_id, timestamp, embedding, decoded, text_ids, latency_ms
        """
        t_start = time.monotonic()

        # 1. Preprocess
        tensor = self.preprocessor.preprocess(frame)  # [1, 3, 384, 384]

        # 2. Encode → embedding
        t_enc = time.monotonic()
        embedding = self.model.get_embedding(tensor)  # [1, embed_dim]
        embed_flat = embedding.squeeze(0)  # [embed_dim]
        t_enc_done = time.monotonic()
        self.metrics.record_encode(t_enc_done - t_enc)

        # 3. Selective decode check
        should_decode = force_decode or self.selective_decoder.should_decode(
            camera_id, embed_flat, timestamp
        )

        result = {
            "camera_id": camera_id,
            "timestamp": timestamp,
            "embedding": embed_flat,
            "decoded": False,
            "text_ids": None,
            "latency_ms": 0.0,
        }

        if should_decode:
            # 4. Run MoE decoder
            t_dec = time.monotonic()
            text_ids = self.model.decoder.generate(
                embedding,
                max_new_tokens=self.max_new_tokens,
                temperature=self.temperature,
            )
            t_dec_done = time.monotonic()
            self.metrics.record_decode(t_dec_done - t_dec)
            self.metrics.decodes_triggered += 1

            if force_decode:
                self.selective_decoder.force_decode(camera_id, embed_flat, timestamp)

            result["decoded"] = True
            result["text_ids"] = text_ids

            # Store latest result
            with self._results_lock:
                self._results[camera_id] = result
        else:
            self.metrics.decodes_skipped += 1

        t_end = time.monotonic()
        result["latency_ms"] = (t_end - t_start) * 1000
        self.metrics.record_total(t_end - t_start)

        return result if should_decode else None

    # ------------------------------------------------------------------
    # Batch inference
    # ------------------------------------------------------------------

    @torch.no_grad()
    def process_batch(
        self,
        camera_ids: list[str],
        frames: list[np.ndarray],
        timestamps: list[float],
    ) -> list[Optional[dict]]:
        """
        Batch inference for multiple cameras.

        Encodes all frames in a single forward pass, then selectively
        decodes only the cameras with semantic shifts.

        Args:
            camera_ids: Camera IDs corresponding to each frame
            frames: List of BGR images
            timestamps: List of timestamps
        Returns:
            List of result dicts (None for cameras that didn't trigger decode)
        """
        if len(frames) == 0:
            return []

        t_start = time.monotonic()

        # 1. Batch preprocess
        batch_tensor = self.preprocessor.preprocess_batch(frames)  # [B, 3, 384, 384]

        # 2. Batch encode
        t_enc = time.monotonic()
        embeddings = self.model.get_embedding(batch_tensor)  # [B, embed_dim]
        t_enc_done = time.monotonic()
        self.metrics.record_encode(t_enc_done - t_enc)

        # 3. Batch selective decode check
        should_decode_list = self.selective_decoder.batch_should_decode(
            camera_ids, embeddings, timestamps
        )

        # 4. Decode only the triggered cameras
        results: list[Optional[dict]] = [None] * len(frames)
        decode_indices = [i for i, sd in enumerate(should_decode_list) if sd]

        if decode_indices:
            # Gather embeddings that need decoding
            decode_embeddings = embeddings[decode_indices]  # [D, embed_dim]

            t_dec = time.monotonic()
            text_ids = self.model.decoder.generate(
                decode_embeddings,
                max_new_tokens=self.max_new_tokens,
                temperature=self.temperature,
            )
            t_dec_done = time.monotonic()
            self.metrics.record_decode(t_dec_done - t_dec)

            for j, idx in enumerate(decode_indices):
                result = {
                    "camera_id": camera_ids[idx],
                    "timestamp": timestamps[idx],
                    "embedding": embeddings[idx],
                    "decoded": True,
                    "text_ids": text_ids[j:j+1] if text_ids is not None else None,
                    "latency_ms": 0.0,
                }
                results[idx] = result
                with self._results_lock:
                    self._results[camera_ids[idx]] = result

        self.metrics.decodes_triggered += len(decode_indices)
        self.metrics.decodes_skipped += len(frames) - len(decode_indices)

        t_end = time.monotonic()
        total_ms = (t_end - t_start) * 1000
        for r in results:
            if r is not None:
                r["latency_ms"] = total_ms
        self.metrics.record_total(t_end - t_start)

        return results

    # ------------------------------------------------------------------
    # Background processing loop
    # ------------------------------------------------------------------

    def _processing_loop(self) -> None:
        """
        Continuously pull keyframes from all cameras and run inference.

        Runs in a background thread, processing all available keyframes
        in batched mode for efficiency.
        """
        logger.info("Edge inference processing loop started")

        while not self._stop_event.is_set():
            try:
                # Collect latest frame from each camera
                all_frames = self.camera_manager.get_all_frames()

                camera_ids = []
                frames = []
                timestamps = []

                for cam_id, frame_data in all_frames.items():
                    if frame_data is None:
                        continue
                    frame, ts = frame_data

                    # Run through keyframe sampler
                    keyframe = self.sampler.process_frame(cam_id, frame, ts)
                    if keyframe is not None:
                        kf_frame, kf_ts = keyframe
                        camera_ids.append(cam_id)
                        frames.append(kf_frame)
                        timestamps.append(kf_ts)

                # Batch inference if we have keyframes
                if frames:
                    self.process_batch(camera_ids, frames, timestamps)

            except Exception as e:
                logger.error(f"Processing loop error: {e}", exc_info=True)

            # Brief sleep to avoid busy-waiting
            self._stop_event.wait(timeout=self._processing_interval)

        logger.info("Edge inference processing loop stopped")

    def start(self) -> None:
        """Start the background processing loop."""
        if self._thread is not None and self._thread.is_alive():
            logger.warning("Processing loop already running")
            return

        self._stop_event.clear()
        self._thread = threading.Thread(
            target=self._processing_loop,
            name="edge-inference-loop",
            daemon=True,
        )
        self._thread.start()

    def stop(self) -> None:
        """Stop the processing loop and all cameras."""
        self._stop_event.set()
        if self._thread is not None:
            self._thread.join(timeout=10.0)
            self._thread = None
        self.camera_manager.stop_all()

    # ------------------------------------------------------------------
    # Results and status
    # ------------------------------------------------------------------

    def get_latest_result(self, camera_id: str) -> Optional[dict]:
        """Get the most recent decode result for a camera."""
        with self._results_lock:
            return self._results.get(camera_id)

    def get_all_results(self) -> dict[str, dict]:
        """Get latest results for all cameras."""
        with self._results_lock:
            return dict(self._results)

    def status(self) -> dict:
        """Full system status: cameras, sampling, inference metrics."""
        return {
            "cameras": self.camera_manager.status(),
            "sampling": self.sampler.stats(),
            "inference": self.metrics.to_dict(),
            "selective_decode": {
                "decode_ratio": round(self.selective_decoder.decode_ratio, 4),
                "compression_ratio": round(self.selective_decoder.compression_ratio, 2),
            },
        }

    # ------------------------------------------------------------------
    # ONNX export helper
    # ------------------------------------------------------------------

    def export_encoder_onnx(
        self,
        output_path: str = "arcisvlm_encoder.onnx",
        opset_version: int = 17,
    ) -> str:
        """
        Export the X-Encoder (ViT) to ONNX for TensorRT optimization.

        The encoder is the main inference bottleneck — exporting to ONNX
        allows conversion to TensorRT FP16/INT8 for 2-4x speedup on Jetson.

        The predictor and decoder remain in PyTorch (they're lightweight
        and MoE routing doesn't map cleanly to ONNX).

        Args:
            output_path: Where to save the ONNX file
            opset_version: ONNX opset version (17 supports all our ops)
        Returns:
            Path to the saved ONNX file
        """
        dummy_input = torch.randn(1, 3, 384, 384, device=self.device)

        encoder = self.model.x_encoder
        encoder.eval()

        torch.onnx.export(
            encoder,
            dummy_input,
            output_path,
            opset_version=opset_version,
            input_names=["image"],
            output_names=["visual_tokens"],
            dynamic_axes={
                "image": {0: "batch_size"},
                "visual_tokens": {0: "batch_size"},
            },
        )

        logger.info(f"Exported encoder to {output_path}")
        return output_path

    def export_predictor_onnx(
        self,
        output_path: str = "arcisvlm_predictor.onnx",
        opset_version: int = 17,
    ) -> str:
        """
        Export the JEPA predictor to ONNX.

        Args:
            output_path: Where to save the ONNX file
            opset_version: ONNX opset version
        Returns:
            Path to the saved ONNX file
        """
        # Predictor takes visual tokens [B, 576, 768] and optional query tokens
        dummy_visual = torch.randn(1, 576, 768, device=self.device)
        # Query IDs: [B, Q] — use a short query for export
        dummy_query = torch.ones(1, 8, dtype=torch.long, device=self.device)
        dummy_mask = torch.ones(1, 8, dtype=torch.bool, device=self.device)

        predictor = self.model.predictor
        predictor.eval()

        torch.onnx.export(
            predictor,
            (dummy_visual, dummy_query, dummy_mask),
            output_path,
            opset_version=opset_version,
            input_names=["visual_tokens", "query_ids", "query_mask"],
            output_names=["embedding"],
            dynamic_axes={
                "visual_tokens": {0: "batch_size"},
                "query_ids": {0: "batch_size", 1: "query_len"},
                "query_mask": {0: "batch_size", 1: "query_len"},
                "embedding": {0: "batch_size"},
            },
        )

        logger.info(f"Exported predictor to {output_path}")
        return output_path


# ---------------------------------------------------------------------------
# Convenience factory
# ---------------------------------------------------------------------------

def create_edge_server(
    config: dict,
    checkpoint_path: Optional[str] = None,
    device: str = "cuda",
) -> EdgeInferenceServer:
    """
    Factory function to create an EdgeInferenceServer from config.

    Args:
        config: Model config dict (same format as training configs)
        checkpoint_path: Path to model checkpoint (.pt file), or None for random init
        device: Target device
    Returns:
        Configured EdgeInferenceServer ready for .add_camera() and .start()
    """
    model = VLJEPAModel(config)

    if checkpoint_path is not None:
        state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
        # Handle common checkpoint formats
        if "model_state_dict" in state_dict:
            state_dict = state_dict["model_state_dict"]
        elif "state_dict" in state_dict:
            state_dict = state_dict["state_dict"]
        model.load_state_dict(state_dict, strict=False)
        logger.info(f"Loaded checkpoint from {checkpoint_path}")

    sc = config.get("selective_decode", {})
    server = EdgeInferenceServer(
        model=model,
        device=device,
        selective_threshold=sc.get("similarity_threshold", 0.95),
        min_decode_interval=sc.get("min_decode_interval", 1.0),
    )

    return server