"""
Inference backends for BADAS.

Supports multiple inference engines:
- PyTorch: Standard PyTorch inference with optional torch.compile
- ONNX: Cross-platform inference via ONNX Runtime
- TensorRT: Maximum performance GPU inference
"""

import numpy as np
import torch
import torch.nn as nn
from abc import ABC, abstractmethod
from typing import List

# Optional imports
try:
    import onnxruntime as ort
    ONNX_AVAILABLE = True
except ImportError:
    ort = None
    ONNX_AVAILABLE = False

try:
    import tensorrt as trt
    TRT_AVAILABLE = True
except ImportError:
    trt = None
    TRT_AVAILABLE = False


def _logits_to_prob(logits: torch.Tensor) -> float:
    """Convert model output logits to a probability scalar.

    Supports both binary (1 logit → sigmoid) and multi-class
    (N logits → softmax, positive class = last index) heads.
    """
    logits = logits.reshape(-1).cpu().float()
    if logits.numel() == 1:
        return torch.sigmoid(logits).item()
    return torch.softmax(logits, dim=0)[-1].item()


def _logits_batch_to_probs(logits: torch.Tensor) -> List[float]:
    """Convert a batch of logits (N, C) or (N,) to a list of N probabilities."""
    logits = logits.cpu().float()
    if logits.dim() == 1 or logits.shape[-1] == 1:
        return torch.sigmoid(logits).reshape(-1).tolist()
    return torch.softmax(logits, dim=-1)[:, -1].tolist()


def _np_logits_to_prob(logits: np.ndarray) -> float:
    """Convert numpy logits to a probability scalar (ONNX / TRT outputs)."""
    logits = logits.reshape(-1).astype(np.float32)
    if logits.size == 1:
        return float(1.0 / (1.0 + np.exp(-logits[0])))
    shifted = logits - logits.max()  # numerically stable softmax
    exp = np.exp(shifted)
    return float(exp[-1] / exp.sum())


class InferenceBackend(ABC):
    """Abstract base class for inference backends."""

    @abstractmethod
    def forward(self, video_tensor: torch.Tensor) -> float:
        """Run forward pass on preprocessed video tensor.

        Args:
            video_tensor: Preprocessed video tensor (1, T, C, H, W)

        Returns:
            Probability value (0-1)
        """
        pass

    def forward_async(self, video_tensor: torch.Tensor):
        """Queue inference without blocking. Default: synchronous forward.

        Subclasses with async GPU support return a GPU tensor (logits)
        without calling synchronize(). Backends without async support
        return a float directly.

        Args:
            video_tensor: Preprocessed video tensor (1, T, C, H, W)

        Returns:
            GPU logits tensor (async) or float (sync fallback)
        """
        return self.forward(video_tensor)

    def collect_result(self, handle) -> float:
        """Extract probability from a forward_async handle.

        Args:
            handle: GPU logits tensor or float returned by forward_async

        Returns:
            Probability value (0-1)
        """
        if isinstance(handle, torch.Tensor):
            return _logits_to_prob(handle)
        return float(handle)

    def forward_batch(self, video_tensor: torch.Tensor) -> List[float]:
        """Batched forward pass. Default: loop over batch dim.

        Args:
            video_tensor: (N, T, C, H, W) preprocessed batch

        Returns:
            List of N probabilities in [0, 1]
        """
        return [self.forward(video_tensor[i:i + 1]) for i in range(video_tensor.shape[0])]


class PyTorchBackend(InferenceBackend):
    """PyTorch inference backend."""

    def __init__(
        self,
        model: nn.Module,
        device: torch.device,
        use_autocast: bool = False
    ):
        self.model = model
        self.device = device
        self.use_autocast = use_autocast

    @staticmethod
    def _cudagraph_mark_step():
        """Signal torch.compile (reduce-overhead) that a new inference step begins."""
        if hasattr(torch, 'compiler') and hasattr(torch.compiler, 'cudagraph_mark_step_begin'):
            torch.compiler.cudagraph_mark_step_begin()

    def forward(self, video_tensor: torch.Tensor) -> float:
        if torch.cuda.is_available():
            self._cudagraph_mark_step()
            if self.use_autocast:
                with torch.no_grad():
                    with torch.amp.autocast('cuda'):
                        logits = self.model(video_tensor)
            else:
                with torch.no_grad():
                    logits = self.model(video_tensor)
            torch.cuda.synchronize()
        else:
            video_tensor = video_tensor.to(self.device)
            with torch.no_grad():
                logits = self.model(video_tensor)

        return _logits_to_prob(logits)

    def forward_async(self, video_tensor: torch.Tensor):
        """Queue model inference without synchronizing (CUDA only).

        Returns a GPU logits tensor. The GPU keeps executing asynchronously.
        Call collect_result() to wait and extract the probability.
        """
        if not torch.cuda.is_available():
            return self.forward(video_tensor)  # CPU: already synchronous

        self._cudagraph_mark_step()
        if self.use_autocast:
            with torch.no_grad():
                with torch.amp.autocast('cuda'):
                    logits = self.model(video_tensor)
        else:
            with torch.no_grad():
                logits = self.model(video_tensor)

        # No synchronize() — GPU is still running, logits tensor not yet ready on CPU
        return logits

    def collect_result(self, handle) -> float:
        """Sync and extract probability. .cpu() waits for the producing CUDA stream."""
        if isinstance(handle, torch.Tensor):
            return _logits_to_prob(handle)
        return float(handle)

    def forward_batch(self, video_tensor: torch.Tensor) -> List[float]:
        """Batched forward pass. (N, T, C, H, W) → N floats."""
        if self.use_autocast:
            with torch.no_grad():
                with torch.amp.autocast('cuda'):
                    logits = self.model(video_tensor)
        else:
            with torch.no_grad():
                logits = self.model(video_tensor)
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        return _logits_batch_to_probs(logits)


class ONNXBackend(InferenceBackend):
    """ONNX Runtime inference backend."""

    def __init__(self, onnx_path: str, use_cuda: bool = True, quiet: bool = False):
        if not ONNX_AVAILABLE:
            raise ImportError("onnxruntime not installed. Install with: pip install onnxruntime-gpu")

        if use_cuda:
            providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
        else:
            providers = ['CPUExecutionProvider']

        if not quiet:
            print(f"   Loading ONNX model: {onnx_path}")

        self.session = ort.InferenceSession(onnx_path, providers=providers)
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name

        # Auto-detect input dtype from model
        input_type = self.session.get_inputs()[0].type
        if 'float16' in input_type:
            self.input_dtype = np.float16
            self.torch_dtype = torch.float16
        else:
            self.input_dtype = np.float32
            self.torch_dtype = torch.float32

        if not quiet:
            print(f"   ONNX provider: {self.session.get_providers()[0]}")
            print(f"   ONNX input dtype: {self.input_dtype.__name__}")

    def forward(self, video_tensor: torch.Tensor) -> float:
        # Convert to numpy with correct dtype for the model
        if self.input_dtype == np.float16:
            video_np = video_tensor.cpu().half().numpy()
        else:
            video_np = video_tensor.cpu().float().numpy()
        outputs = self.session.run([self.output_name], {self.input_name: video_np})
        logits = outputs[0]
        return _np_logits_to_prob(logits)


class TensorRTBackend(InferenceBackend):
    """
    TensorRT inference backend.  This is the canonical TRT backend and the
    recommended replacement for the deprecated ``ZeroCopyTensorRTEngine``.

    Implements the full ``InferenceBackend`` protocol (``forward``,
    ``forward_async``, ``collect_result``, ``forward_batch``) and provides
    additional zero-copy entry points for maximum GPU performance.

    Inference modes
    ---------------
    ``forward(tensor)``
        Standard mode.  Copies the input to a pre-allocated internal buffer,
        then runs TensorRT asynchronously on a dedicated CUDA stream.
        Accepts any dtype or device; handles conversion internally.
        Use this when correctness and simplicity matter more than the last
        few microseconds of copy overhead.

    ``forward_zero_copy(tensor)``
        Zero-copy mode.  Points TensorRT directly at the caller's CUDA tensor
        without any copy.  The tensor must already be on CUDA, contiguous, and
        FP32.  Use this when the tensor comes from ``GPUPreprocessor`` or
        another GPU pipeline stage and you want to eliminate the copy.
        Equivalent to the deprecated ``ZeroCopyTensorRTEngine.infer_from_tensor``,
        but with correct sigmoid output and proper CUDA stream ordering.

    ``forward_from_ptr(ptr)``
        Maximum zero-copy mode.  Accepts a raw GPU memory pointer (e.g. from
        ``GPUPreprocessor.get_data_ptr()``).  No tensor operations at all.
        Equivalent to the deprecated ``ZeroCopyTensorRTEngine.infer_from_ptr``,
        but with correct sigmoid output and proper CUDA stream ordering.

    ``forward_async(tensor)`` / ``collect_result(handle)``
        Async mode.  Submits TRT work without blocking Python; returns a
        ``torch.cuda.Event``.  Call ``collect_result(event)`` later to wait
        and extract the probability.

    CUDA stream
    -----------
    All execution runs on a dedicated ``_trt_stream`` (not the default stream
    0).  This avoids the extra ``cudaStreamSynchronize`` calls that TensorRT
    inserts when running on stream 0, and makes async overlap with other GPU
    work correct.

    FP16 support
    ------------
    Always use FP32 for input/output bindings.  TensorRT FP16 engines
    (built with ``--fp16``) use FP16 for internal computations but keep
    FP32 bindings.  The engine handles the conversion automatically.

    Migration from ZeroCopyTensorRTEngine
    --------------------------------------
    Replace::

        engine = ZeroCopyTensorRTEngine("model.trt")
        logit = engine.infer_from_tensor(tensor)   # WARNING: returns raw logit
        logit = engine.infer_from_ptr(ptr)          # WARNING: returns raw logit

    With::

        backend = TensorRTBackend("model.trt")
        prob = backend.forward_zero_copy(tensor)   # returns probability in [0, 1]
        prob = backend.forward_from_ptr(ptr)        # returns probability in [0, 1]
        prob = backend.forward(tensor)              # standard path (buffer copy)
    """

    def __init__(self, engine_path: str, quiet: bool = False):
        if not TRT_AVAILABLE:
            raise ImportError("tensorrt not installed. Install with: pip install tensorrt")

        if not quiet:
            print(f"   Loading TensorRT engine: {engine_path}")

        self.logger = trt.Logger(trt.Logger.WARNING)
        self.runtime = trt.Runtime(self.logger)

        with open(engine_path, 'rb') as f:
            self.engine = self.runtime.deserialize_cuda_engine(f.read())

        if self.engine is None:
            raise RuntimeError(f"Failed to load TensorRT engine: {engine_path}")

        self.context = self.engine.create_execution_context()
        self._setup_io(trt, quiet)

    def _setup_io(self, trt, quiet: bool):
        """Setup input/output tensor info and buffers."""
        self.input_name = None
        self.output_name = None

        for i in range(self.engine.num_io_tensors):
            name = self.engine.get_tensor_name(i)
            shape = self.engine.get_tensor_shape(name)
            mode = self.engine.get_tensor_mode(name)

            if mode == trt.TensorIOMode.INPUT:
                self.input_name = name
                self.input_shape = tuple(shape)
                if not quiet:
                    print(f"   TRT Input: {name} {list(shape)}")
            else:
                self.output_name = name
                self.output_shape = tuple(shape) if shape[0] > 0 else (1,)
                if not quiet:
                    print(f"   TRT Output: {name} {list(shape)}")

        # Set input shape for dynamic shapes
        self.context.set_input_shape(self.input_name, self.input_shape)

        # Pre-allocate GPU buffers
        # IMPORTANT: Always use FP32 for input/output bindings.
        # TensorRT FP16 engines (built with --fp16) use FP16 for INTERNAL
        # computations, but the input/output bindings remain FP32.
        # The engine automatically converts inputs to FP16 internally.
        self.dtype = torch.float32

        self.input_buffer = torch.empty(
            self.input_shape,
            dtype=torch.float32,
            device='cuda'
        ).contiguous()

        # Output is always FP32
        self.output_buffer = torch.empty(
            self.output_shape,
            dtype=torch.float32,
            device='cuda'
        ).contiguous()

        # Dedicated CUDA stream: avoids extra cudaStreamSynchronize() calls that
        # TRT inserts when running on the default stream (stream 0).
        self._trt_stream = torch.cuda.Stream()

        if not quiet:
            print(f"   TRT Precision: FP32 input/output (engine may use FP16 internally)")

    def forward(self, video_tensor: torch.Tensor) -> float:
        """
        Standard forward pass (copies input to internal buffer).

        Args:
            video_tensor: Input tensor (any dtype, will be converted to FP32)

        Returns:
            Probability value (0-1)
        """
        # Always use FP32 for TRT input (engine handles internal conversion)
        if video_tensor.dtype != torch.float32:
            video_tensor = video_tensor.float()

        if not video_tensor.is_cuda:
            video_tensor = video_tensor.cuda()
        if not video_tensor.is_contiguous():
            video_tensor = video_tensor.contiguous()

        # Run on dedicated stream so TRT doesn't insert extra syncs on stream 0.
        # wait_stream ensures TRT sees any default-stream work (e.g. preprocessing).
        self._trt_stream.wait_stream(torch.cuda.current_stream())
        with torch.cuda.stream(self._trt_stream):
            self.input_buffer.copy_(video_tensor, non_blocking=True)
            self.context.set_tensor_address(self.input_name, self.input_buffer.data_ptr())
            self.context.set_tensor_address(self.output_name, self.output_buffer.data_ptr())
            self.context.execute_async_v3(self._trt_stream.cuda_stream)
        self._trt_stream.synchronize()

        logits = float(self.output_buffer.item())
        return 1.0 / (1.0 + np.exp(-logits))

    def forward_async(self, video_tensor: torch.Tensor):
        """Submit TRT inference without blocking. Returns a CUDA Event handle.

        Python returns after ~0.1ms; GPU runs for ~24ms asynchronously.
        Call collect_result(handle) to wait and extract the probability.

        Args:
            video_tensor: Input tensor (any dtype, will be converted to FP32)

        Returns:
            torch.cuda.Event that signals when TRT is done
        """
        if video_tensor.dtype != torch.float32:
            video_tensor = video_tensor.float()

        if not video_tensor.is_cuda:
            video_tensor = video_tensor.cuda()
        if not video_tensor.is_contiguous():
            video_tensor = video_tensor.contiguous()

        event = torch.cuda.Event()
        self._trt_stream.wait_stream(torch.cuda.current_stream())
        with torch.cuda.stream(self._trt_stream):
            self.input_buffer.copy_(video_tensor, non_blocking=True)
            self.context.set_tensor_address(self.input_name, self.input_buffer.data_ptr())
            self.context.set_tensor_address(self.output_name, self.output_buffer.data_ptr())
            self.context.execute_async_v3(self._trt_stream.cuda_stream)
            event.record(self._trt_stream)
        return event  # GPU still running; Python returns here

    def collect_result(self, handle) -> float:
        """Wait for async TRT inference to complete and extract probability."""
        if isinstance(handle, torch.cuda.Event):
            handle.synchronize()
            logits = float(self.output_buffer.item())
            return 1.0 / (1.0 + np.exp(-logits))
        return float(handle)

    def forward_zero_copy(self, video_tensor: torch.Tensor) -> float:
        """
        Zero-copy forward pass (uses input tensor directly).

        This avoids memory copies for maximum performance.
        Tensor must be:
        - On CUDA device
        - Contiguous
        - FP32 dtype

        Args:
            video_tensor: Input tensor from GPUPreprocessor.get_input_tensor()

        Returns:
            Probability value (0-1)
        """
        # Ensure tensor is ready for zero-copy
        if not video_tensor.is_cuda:
            raise ValueError("Tensor must be on CUDA device for zero-copy")

        # Always use FP32 for TRT input (engine handles internal conversion)
        if video_tensor.dtype != torch.float32:
            video_tensor = video_tensor.float()

        if not video_tensor.is_contiguous():
            video_tensor = video_tensor.contiguous()

        # Set tensor addresses directly (zero-copy)
        self._trt_stream.wait_stream(torch.cuda.current_stream())
        self.context.set_tensor_address(self.input_name, video_tensor.data_ptr())
        self.context.set_tensor_address(self.output_name, self.output_buffer.data_ptr())
        self.context.execute_async_v3(self._trt_stream.cuda_stream)
        self._trt_stream.synchronize()

        logits = float(self.output_buffer.item())
        return 1.0 / (1.0 + np.exp(-logits))

    def forward_from_ptr(self, input_ptr: int) -> float:
        """
        Forward pass from raw GPU memory pointer.

        Maximum zero-copy mode - no tensor operations at all.
        NOTE: The data at input_ptr must be in the correct dtype (FP16 or FP32).

        Args:
            input_ptr: GPU memory pointer from GPUPreprocessor.get_data_ptr()

        Returns:
            Probability value (0-1)
        """
        self.context.set_tensor_address(self.input_name, input_ptr)
        self.context.set_tensor_address(self.output_name, self.output_buffer.data_ptr())
        self.context.execute_async_v3(self._trt_stream.cuda_stream)
        self._trt_stream.synchronize()

        logits = float(self.output_buffer.item())
        return 1.0 / (1.0 + np.exp(-logits))


class CosmosVLMBackend(InferenceBackend):
    """
    vLLM-based backend for Qwen3-VL family models (Cosmos-Reason2, Qwen3-VL).

    Accepts raw BGR frames (not preprocessed tensors) — set needs_raw_frames=True
    so BADAS bypasses VideoPreprocessor and passes the raw rolling buffer instead.

    Scores a window by extracting P(A) / (P(A) + P(B)) from the first-token
    logprobs — same A/B logit approach used for the Kaggle benchmark.

    Supports:
        nvidia/Cosmos-Reason2-2B
        Qwen/Qwen3-VL-2B-Instruct  (and larger variants)
        Any HF safetensors directory with the same architecture
    """

    needs_raw_frames: bool = True

    SYSTEM_PROMPT = (
        "You are an expert in driving safety. "
        "Analyze dashcam footage and assess collision risk."
    )
    USER_PROMPT = (
        "This is a 2-second dashcam clip from a moving vehicle.\n\n"
        "Will the ego vehicle be involved in a collision within the next few seconds?\n\n"
        "A: Yes — a collision is imminent or already occurring.\n"
        "B: No — the driving situation appears safe.\n\n"
        "Answer with a single letter."
    )

    def __init__(
        self,
        model_path: str,
        fps: float = 4.0,
        max_pixels: int = 256 * 256,
        gpu_memory_utilization: float = 0.85,
        quiet: bool = False,
    ):
        import os
        import torch
        from transformers import AutoProcessor
        from vllm import LLM, SamplingParams

        self.fps = fps
        self.max_pixels = max_pixels

        os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")

        if not quiet:
            print(f"   Loading VLM backend: {model_path}")

        self.processor = AutoProcessor.from_pretrained(model_path)
        tok = self.processor.tokenizer
        self.tok_A = tok.encode("A", add_special_tokens=False)[0]
        self.tok_B = tok.encode("B", add_special_tokens=False)[0]

        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            # mem_get_info() reports physical free memory across ALL processes,
            # unlike memory_allocated() which only counts this process's usage.
            free, total = torch.cuda.mem_get_info(0)
            free_fraction = free / total
            # Leave a 3% safety margin; never ask for more than available
            if gpu_memory_utilization > free_fraction - 0.03:
                adjusted = round(max(0.1, free_fraction - 0.05), 2)
                if not quiet:
                    print(f"   Adjusting gpu_memory_utilization: "
                          f"{gpu_memory_utilization:.2f} → {adjusted:.2f} "
                          f"(free: {free_fraction:.2f})")
                gpu_memory_utilization = adjusted

        self.llm = LLM(
            model=model_path,
            max_model_len=4096,
            dtype="float16",
            gpu_memory_utilization=gpu_memory_utilization,
            limit_mm_per_prompt={"video": 1},
            disable_log_stats=True,
        )
        self.sampling_params = SamplingParams(max_tokens=1, temperature=0, logprobs=20)

        if not quiet:
            print(f"   VLM ready | A={self.tok_A}, B={self.tok_B}")

    def _resize_frame(self, frame_bgr: np.ndarray):
        """BGR numpy → resized PIL RGB, capped at max_pixels."""
        import cv2
        from PIL import Image
        arr = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
        h, w = arr.shape[:2]
        if h * w > self.max_pixels:
            scale = (self.max_pixels / (h * w)) ** 0.5
            new_w = max(16, round(w * scale / 16) * 16)
            new_h = max(16, round(h * scale / 16) * 16)
            arr = cv2.resize(arr, (new_w, new_h), interpolation=cv2.INTER_AREA)
        return Image.fromarray(arr)

    def forward(self, frames: List[np.ndarray]) -> float:
        """Score a window of BGR frames. Returns P(collision) in [0, 1]."""
        import math
        pil_frames = [self._resize_frame(f) for f in frames]
        msgs = [
            {"role": "system", "content": self.SYSTEM_PROMPT},
            {"role": "user", "content": [
                {"type": "video", "video": pil_frames, "fps": self.fps},
                {"type": "text",  "text": self.USER_PROMPT},
            ]},
        ]
        text = self.processor.apply_chat_template(
            msgs, tokenize=False, add_generation_prompt=True)
        frames_np = np.stack([np.array(f) for f in pil_frames])
        n = len(pil_frames)
        metadata = {"total_num_frames": n, "fps": self.fps, "frames_indices": list(range(n))}
        outputs = self.llm.generate(
            [{"prompt": text, "multi_modal_data": {"video": (frames_np, metadata)}}],
            self.sampling_params,
            use_tqdm=False,
        )
        lp = outputs[0].outputs[0].logprobs[0]
        la, lb = lp.get(self.tok_A), lp.get(self.tok_B)
        if la is not None and lb is not None:
            pa, pb = math.exp(la.logprob), math.exp(lb.logprob)
            return pa / (pa + pb)
        return 0.5

    def forward_async(self, frames: List[np.ndarray]) -> float:
        """vLLM is internally async but externally synchronous; returns float directly."""
        return self.forward(frames)


class GeminiVLMBackend(InferenceBackend):
    """
    Gemini 2.5 Pro backend for collision probability estimation.

    Since Gemini 2.5 Pro does not support logprobs, the model is prompted to
    output a single decimal probability in [0.0, 1.0] directly.

    Supports two auth modes:
    - API key (env GEMINI_API_KEY): for vanilla gemini-2.5-pro via Google AI API
    - Vertex AI (gcloud token): for fine-tuned endpoints (model_id starts with "projects/")

    Sends raw BGR frames as inline base64 JPEG images.
    """

    needs_raw_frames: bool = True

    SYSTEM_PROMPT = (
        "You are an expert in driving safety analyzing dashcam footage."
    )
    USER_PROMPT = (
        "These are frames sampled from a 2-second dashcam clip.\n\n"
        "Will the ego vehicle be involved in a collision within the next few seconds?\n\n"
        "Respond with ONLY a decimal number from 0.0 to 1.0, where:\n"
        "0.0 = completely safe, no collision risk\n"
        "1.0 = certain collision is imminent or already occurring\n\n"
        "Output ONLY the number, nothing else."
    )

    def __init__(
        self,
        model_id: str,
        api_key: str | None = None,
        fps: float = 4.0,
        max_image_size: int = 512,
        quiet: bool = False,
        temperatures: List[float] | None = None,
        max_workers: int = 6,
    ):
        import os
        import threading

        self.model_id = model_id
        self.fps = fps
        self.max_image_size = max_image_size
        self.quiet = quiet
        self._temperatures: List[float] = temperatures if temperatures is not None else [0.0, 0.3, 0.7]
        self._max_workers = max(max_workers, len(self._temperatures))
        # Persistent executor shared across all forward() calls (avoids per-call spawn overhead
        # and allows deep window-level parallelism from the outer caller).
        from concurrent.futures import ThreadPoolExecutor as _TPE
        self._executor = _TPE(max_workers=self._max_workers)

        self._is_vertex = model_id.startswith("projects/")
        self._token_lock = threading.Lock()

        if self._is_vertex:
            self._token: str | None = None
            self._token_expiry: float = 0.0
            if not quiet:
                print(f"   Gemini backend (Vertex AI): {model_id}")
        else:
            self._api_key = api_key or os.environ.get("GEMINI_API_KEY")
            if not self._api_key:
                raise ValueError(
                    "GEMINI_API_KEY env var or api_key parameter required for "
                    "non-Vertex Gemini models"
                )
            if not quiet:
                print(f"   Gemini backend (API key): {model_id}")

    def _get_vertex_token(self) -> str:
        """Return a valid gcloud access token, refreshing every 30 min. Thread-safe."""
        import subprocess
        import time

        with self._token_lock:
            if time.time() > self._token_expiry:
                result = subprocess.run(
                    ["gcloud", "auth", "print-access-token"],
                    capture_output=True, text=True, check=True,
                )
                self._token = result.stdout.strip()
                self._token_expiry = time.time() + 1800  # 30 min
            return self._token

    def _encode_frame(self, frame_bgr: np.ndarray) -> str:
        """BGR numpy frame → base64-encoded JPEG string."""
        import base64
        import cv2

        h, w = frame_bgr.shape[:2]
        max_dim = max(h, w)
        if max_dim > self.max_image_size:
            scale = self.max_image_size / max_dim
            new_w = max(1, round(w * scale))
            new_h = max(1, round(h * scale))
            frame_bgr = cv2.resize(frame_bgr, (new_w, new_h), interpolation=cv2.INTER_AREA)
        _, buf = cv2.imencode(".jpg", frame_bgr, [cv2.IMWRITE_JPEG_QUALITY, 85])
        return base64.b64encode(buf).decode("utf-8")

    def _call_api(self, frames: List[np.ndarray], temperature: float = 0.0) -> str:
        """Send frames to Gemini and return raw response text. Retries on rate-limit errors."""
        import time
        import requests

        parts = []
        for frame in frames:
            parts.append({
                "inlineData": {
                    "mimeType": "image/jpeg",
                    "data": self._encode_frame(frame),
                }
            })
        parts.append({"text": self.USER_PROMPT})

        payload = {
            "systemInstruction": {"parts": [{"text": self.SYSTEM_PROMPT}]},
            "contents": [{"role": "user", "parts": parts}],
            "generationConfig": {"temperature": temperature, "maxOutputTokens": 4096},
        }

        if self._is_vertex:
            url = (
                f"https://us-central1-aiplatform.googleapis.com/v1/"
                f"{self.model_id}:generateContent"
            )
        else:
            url = (
                f"https://generativelanguage.googleapis.com/v1beta/models/"
                f"{self.model_id}:generateContent?key={self._api_key}"
            )

        max_retries = 8
        backoff = 5.0  # initial backoff in seconds
        for attempt in range(max_retries):
            if self._is_vertex:
                token = self._get_vertex_token()
                headers = {
                    "Authorization": f"Bearer {token}",
                    "Content-Type": "application/json",
                }
            else:
                headers = {"Content-Type": "application/json"}

            try:
                resp = requests.post(url, json=payload, headers=headers, timeout=120)
            except requests.exceptions.RequestException as e:
                if attempt < max_retries - 1:
                    if not self.quiet:
                        print(f"   Gemini request error (attempt {attempt+1}): {e} — retrying in {backoff:.0f}s", flush=True)
                    time.sleep(backoff)
                    backoff = min(backoff * 2, 120)
                    continue
                raise

            if resp.status_code == 401 and self._is_vertex:
                # Token may be stale — force refresh and retry immediately
                self._token_expiry = 0.0
                if not self.quiet:
                    print(f"   Gemini 401 (attempt {attempt+1}) — forcing token refresh", flush=True)
                continue

            if resp.status_code == 429 or resp.status_code >= 500:
                if attempt < max_retries - 1:
                    retry_after = float(resp.headers.get("Retry-After", backoff))
                    wait = max(retry_after, backoff)
                    if not self.quiet:
                        print(f"   Gemini HTTP {resp.status_code} (attempt {attempt+1}) — retrying in {wait:.0f}s", flush=True)
                    time.sleep(wait)
                    backoff = min(backoff * 2, 120)
                    continue
                resp.raise_for_status()

            resp.raise_for_status()
            data = resp.json()
            try:
                return data["candidates"][0]["content"]["parts"][0]["text"]
            except (KeyError, IndexError) as e:
                raise RuntimeError(f"Unexpected Gemini response: {data}") from e

        raise RuntimeError(f"Gemini API failed after {max_retries} attempts")

    def _score_at_temperature(self, frames: List[np.ndarray], temperature: float) -> float:
        """Call API at a single temperature and parse the probability."""
        import re
        text = self._call_api(frames, temperature).strip()
        match = re.search(r'\d+\.?\d*', text)
        return max(0.0, min(1.0, float(match.group()))) if match else 0.5

    def forward(self, frames: List[np.ndarray]) -> float:
        """Score a window of BGR frames. Returns P(collision) in [0, 1].

        If multiple temperatures are configured, all calls are dispatched to the
        persistent thread pool and probabilities are averaged for better calibration.
        This method is thread-safe and can be called concurrently from outer pools.
        """
        from concurrent.futures import as_completed

        if len(self._temperatures) == 1:
            return self._score_at_temperature(frames, self._temperatures[0])

        futures = {
            self._executor.submit(self._score_at_temperature, frames, t): t
            for t in self._temperatures
        }
        probs: List[float] = []
        for fut in as_completed(futures):
            try:
                probs.append(fut.result())
            except Exception as e:
                if not self.quiet:
                    print(f"   Gemini temp={futures[fut]:.1f} error: {e}", flush=True)

        return sum(probs) / len(probs) if probs else 0.5

    def forward_async(self, frames: List[np.ndarray]) -> float:
        """Gemini API calls are handled via thread pool internally."""
        return self.forward(frames)


__all__ = [
    'InferenceBackend',
    'PyTorchBackend',
    'ONNXBackend',
    'TensorRTBackend',
    'CosmosVLMBackend',
    'GeminiVLMBackend',
    'ONNX_AVAILABLE',
    'TRT_AVAILABLE',
]
