"""
HazardDescriptor — local VLM scene description for BADAS attribution events.

Uses Qwen/Qwen3-VL-2B-Instruct via vLLM for fast inference (~120ms per call
after warm-up, vs ~1.8s with plain HuggingFace generate).
"""

import cv2
import numpy as np
import torch
from typing import List, Optional, Tuple


_SYSTEM_PROMPT = (
    "You are a real-time driving safety assistant. "
    "A dashcam frame has been flagged for collision risk. "
    "The action should be a short driver command that would prevent the danger, "
    "choosing the minimal maneuver needed (e.g., brake, brake immediately, slow down, steer away, pay attention, stay in your lane). "
    "CRITICAL spatial rule: if the hazard is on one side of the frame, the escape direction is toward the opposite side — always steer away from the marked hazard, never toward it. "
    "If there is no safe lateral escape (e.g., wall, traffic, narrow road), the correct action is to brake. "
    'Respond only with valid JSON: {"reasoning": "<one sentence>", "action": "<2–4 word command>"}'
)

_USER_TEMPLATE = (
    "Collision risk {risk:.0%}. "
    "A hazard is marked with a red box on the {side} side of the frame. "
    "What is the danger and what should the driver do to avoid it?"
)

_FALLBACK_TEMPLATE = (
    "Collision risk {risk:.0%}. "
    "The scene is dangerous. Why is it risky and what should the driver do?"
)


def _parse_vlm_output(text: str) -> str:
    """Parse JSON model output into 'Action: ...\nReasoning: ...' display string.

    Tries to extract ``reasoning`` and ``action`` fields from JSON.
    Falls back to returning *text* unchanged if parsing fails.
    """
    import json
    import re

    # Strip markdown code fences if the model wraps its JSON
    clean = re.sub(r"```(?:json)?\s*|\s*```", "", text).strip()

    # Some models emit trailing content after the closing brace; extract first {...}
    m = re.search(r"\{.*?\}", clean, re.DOTALL)
    if m:
        clean = m.group(0)

    try:
        data = json.loads(clean)
        reasoning = str(data.get("reasoning", "")).strip()
        action    = str(data.get("action", "")).strip()
        if action and reasoning:
            return f"Action: {action}\nReasoning: {reasoning}"
        if action:
            return f"Action: {action}"
    except (json.JSONDecodeError, ValueError):
        pass

    return text  # fallback: return raw output


class HazardDescriptor:
    """
    Local VLM wrapper for generating driver hazard descriptions.
    Uses vLLM for fast inference with CUDA graph acceleration.

    Usage:
        hd = HazardDescriptor()
        hd.load()
        text = hd.describe(frame_bgr, bbox=(x1,y1,x2,y2), risk=0.82)
    """

    def __init__(
        self,
        model_id: str = "nexar-ai/badas-reason-qwen3-vl-4b",
        device: Optional[str] = None,
        max_new_tokens: int = 150,
        max_pixels: int = 256 * 256,
        gpu_memory_utilization: float = 0.15,
    ):
        self.model_id = model_id
        # device is stored for API consistency; vLLM auto-selects the GPU device
        # and does not accept an explicit device argument.
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.max_new_tokens = max_new_tokens
        self.max_pixels = max_pixels
        self.gpu_memory_utilization = gpu_memory_utilization
        self._llm = None
        self._processor = None
        self._sampling_params = None

    @property
    def is_loaded(self) -> bool:
        return self._llm is not None and self._sampling_params is not None

    def load(self) -> None:
        """Load model via vLLM, capture CUDA graphs, and warm up. Safe to call multiple times."""
        if self.is_loaded:
            return
        import os
        # vLLM spawns an EngineCore subprocess; on Linux the default fork start
        # method can't re-initialise CUDA if the parent already has a CUDA
        # context.  Force 'spawn' so the subprocess starts clean.
        os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
        from vllm import LLM, SamplingParams
        from transformers import AutoProcessor

        self._processor = AutoProcessor.from_pretrained(
            self.model_id, trust_remote_code=True
        )
        # Drain PyTorch's caching allocator before vLLM takes its initial GPU
        # memory snapshot.  If torch.compile has recently freed cached blocks
        # they stay in the allocator unless we flush them; vLLM's profiler then
        # sees memory *increase* during its own dummy forward pass and hits an
        # assertion ("current free > initial free").
        if torch.cuda.is_available():
            torch.cuda.synchronize()
            torch.cuda.empty_cache()
        self._llm = LLM(
            model=self.model_id,
            max_model_len=2048,
            dtype="float16",
            gpu_memory_utilization=self.gpu_memory_utilization,
            disable_log_stats=True,
        )
        self._sampling_params = SamplingParams(
            max_tokens=self.max_new_tokens,
            temperature=0,
            skip_special_tokens=True,
        )

        # Warm up: trigger JIT compilation of the vision encoder so the first
        # real describe() call is fast.  Use a frame sized to match max_pixels.
        side = int(self.max_pixels ** 0.5)
        dummy = np.zeros((side, side, 3), dtype=np.uint8)
        self.describe(dummy, bbox=(0, 0, side // 2, side // 2), risk=0.5)

    def describe(
        self,
        frame_bgr: np.ndarray,
        bboxes: Optional[List[Tuple[int, int, int, int, float]]] = None,
        risk: float = 0.0,
        img_size: int = 256,
        describe_if_no_focus: bool = True,
        # backward-compat: accept a single bare bbox tuple
        bbox: Optional[Tuple[int, int, int, int]] = None,
    ) -> Optional[str]:
        """Generate a one-sentence hazard description for the driver.

        Args:
            frame_bgr:            Raw BGR frame from video/camera (any resolution).
            bboxes:               Ranked list of ``(x1, y1, x2, y2, importance)``
                                  from :meth:`SpatialAttribution.get_bboxes`.
                                  Pass ``[]`` when attention is scattered.
            risk:                 Smoothed risk probability in [0, 1].
            img_size:             Side length (px) of the attribution model's input
                                  space used when computing bboxes.  Default 256.
            describe_if_no_focus: When ``bboxes`` is empty (scattered attention),
                                  still generate a description using the full frame
                                  with a prompt that acknowledges the uncertainty.
                                  Set to ``False`` to return ``None`` instead.
            bbox:                 Deprecated single-bbox fallback for backward
                                  compatibility.  Ignored when ``bboxes`` is given.

        Returns:
            Two-line string in ``"Action: ...\\nReasoning: ..."`` format,
            or ``None`` if model not loaded or ``describe_if_no_focus=False``
            with no focal region.
        """
        if not self.is_loaded:
            return None

        # ── backward-compat shim ──────────────────────────────────────────────
        if bboxes is None:
            if bbox is not None:
                bboxes = [(*bbox, 1.0)]  # wrap legacy tuple, fake importance=1
            else:
                bboxes = []

        # ── scattered / no focus ──────────────────────────────────────────────
        if not bboxes:
            if not describe_if_no_focus:
                return None
            img    = cv2.cvtColor(frame_bgr.copy(), cv2.COLOR_BGR2RGB)
            prompt = _FALLBACK_TEMPLATE.format(risk=risk)
            result = self._run_vlm(img, prompt)
            return _parse_vlm_output(result) if self.is_loaded else None

        # ── one or more focal regions ─────────────────────────────────────────
        img = self._prepare_image(frame_bgr, bboxes, model_input_size=img_size)
        # Compute horizontal position of primary hazard bbox so the VLM knows
        # which side to steer away from (right hazard → steer left, etc.).
        x1, y1, x2, y2, _ = bboxes[0]
        cx = (x1 + x2) / 2 / img_size
        if cx < 0.38:
            side = "left"
        elif cx > 0.62:
            side = "right"
        else:
            side = "center"
        prompt = _USER_TEMPLATE.format(risk=risk, side=side)
        result = self._run_vlm(img, prompt)
        return _parse_vlm_output(result) if self.is_loaded else None

    def _prepare_image(
        self,
        frame_bgr: np.ndarray,
        bboxes: List[Tuple[int, int, int, int, float]],
        model_input_size: int = 256,
    ) -> np.ndarray:
        """Return an RGB copy of the frame with all bboxes drawn.

        Boxes are expected in model-input pixel space (default 256×256) and are
        scaled to the actual frame resolution before drawing.  The primary box
        (highest importance) is drawn with full red; secondary boxes are drawn
        with a thinner, dimmer outline so the priority ordering is visually clear.
        """
        out = cv2.cvtColor(frame_bgr.copy(), cv2.COLOR_BGR2RGB)
        H, W = out.shape[:2]
        sx = W / model_input_size
        sy = H / model_input_size
        base_thickness = max(2, round(W / 320))

        for rank, (x1, y1, x2, y2, importance) in enumerate(bboxes):
            x1 = int(x1 * sx);  y1 = int(y1 * sy)
            x2 = int(x2 * sx);  y2 = int(y2 * sy)
            # Primary box: full red, full thickness.
            # Secondary boxes: progressively dimmer and thinner.
            intensity = max(80, int(220 * (importance / bboxes[0][4])))
            thickness = max(1, base_thickness - rank)
            cv2.rectangle(out, (x1, y1), (x2, y2), (intensity, 30, 30), thickness)
        return out

    def _run_vlm(self, image_rgb: np.ndarray, user_text: str) -> str:
        """Build the chat payload, run via vLLM, return the decoded text."""
        from PIL import Image

        # Cap vision token count: a 1280×720 frame would produce ~900 vision
        # tokens; capping at max_pixels (default 256²=65536) brings it down to
        # ~63 tokens, giving ~15× faster vision encoding.
        h, w = image_rgb.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)
            image_rgb = cv2.resize(image_rgb, (new_w, new_h), interpolation=cv2.INTER_AREA)

        pil_img = Image.fromarray(image_rgb)

        messages = [
            {"role": "system", "content": _SYSTEM_PROMPT},
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": pil_img},
                    {"type": "text", "text": user_text},
                ],
            },
        ]

        text = self._processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )

        try:
            outputs = self._llm.generate(
                {"prompt": text, "multi_modal_data": {"image": pil_img}},
                self._sampling_params,
                use_tqdm=False,
            )
        except Exception as e:
            # Engine subprocess died (OOM, CUDA fault, etc.). Mark dead so
            # is_loaded returns False and future describe() calls skip the VLM.
            print(f"[HazardDescriptor] VLM engine error — disabling: {e}", flush=True)
            self._llm = None
            self._sampling_params = None
            return ""
        return outputs[0].outputs[0].text.strip()
