"""
Vision-Language reasoning engine, powered by Qwen2.5-VL-3B-Instruct.

This is the "brain" of the app: it doesn't just caption the image, it is
prompted to reason like a triage officer and return a structured JSON
judgment (risk score + recommended actions), grounded with YOLO's object
counts to reduce hallucination.
"""

import json
import re
import os
from pathlib import Path

import torch
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info

from .config import VLM_MODEL_ID, DEVICE, DTYPE, VLM_MAX_NEW_TOKENS

PROMPT_PATH = Path(__file__).parent.parent / "prompts" / "triage_prompt.txt"


class DisasterVLM:
    def __init__(self, model_id: str = VLM_MODEL_ID):
        self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
            model_id,
            torch_dtype=DTYPE,
            device_map="auto" if DEVICE == "cuda" else None,
        )
        if DEVICE != "cuda":
            self.model.to(DEVICE)
        self.processor = AutoProcessor.from_pretrained(model_id)
        self.prompt_template = PROMPT_PATH.read_text()

    def _build_prompt(self, yolo_context: str) -> str:
        return self.prompt_template.format(yolo_context=yolo_context)

    def analyze(self, image_path: str, yolo_context: str = ""):
        """Returns (parsed_dict, raw_text). parsed_dict is {} if JSON parsing
        fails completely (caller falls back to safe defaults)."""
        prompt_text = self._build_prompt(yolo_context)

        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image_path},
                    {"type": "text", "text": prompt_text},
                ],
            }
        ]

        chat_prompt = self.processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        image_inputs, video_inputs = process_vision_info(messages)
        inputs = self.processor(
            text=[chat_prompt],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        ).to(self.model.device)

        with torch.no_grad():
            generated_ids = self.model.generate(
                **inputs, max_new_tokens=VLM_MAX_NEW_TOKENS, do_sample=False
            )
        trimmed = [
            out_ids[len(in_ids):]
            for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
        ]
        raw_text = self.processor.batch_decode(
            trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
        )[0]

        return self._safe_parse_json(raw_text), raw_text

    @staticmethod
    def _safe_parse_json(text: str) -> dict:
        """Models occasionally wrap JSON in ```json fences or add stray text.
        Strip fences first, then fall back to extracting the first {...} block."""
        cleaned = text.strip()
        cleaned = re.sub(r"^```(json)?", "", cleaned).strip()
        cleaned = re.sub(r"```$", "", cleaned).strip()
        try:
            return json.loads(cleaned)
        except json.JSONDecodeError:
            pass

        match = re.search(r"\{.*\}", cleaned, re.DOTALL)
        if match:
            try:
                return json.loads(match.group(0))
            except json.JSONDecodeError:
                pass
        return {}