"""
MLRS: InternVL3.5 (LoRA) reasons about the image and emits a segmentation tool call

    <think>...</think><segmentation>{"noun phrase": "...", "objects": [{"bbox": [x1, y1, x2, y2]}]}</segmentation>

(boxes on a 0-999 grid). For every parsed noun phrase, SAM3 (finetuned mask decoder)
is prompted with the phrase and its boxes and returns one semantic mask, restricted
to the predicted boxes. Non-segmentation tasks are answered with `<answer>...</answer>`
and produce no masks.
"""
from __future__ import annotations

import json
import logging
import os
import re
from dataclasses import dataclass, field
from math import isfinite
from typing import Any, Dict, List, Optional, Sequence, Tuple

import torch
import torch.nn.functional as F
from huggingface_hub import snapshot_download
from peft import PeftModel
from safetensors.torch import load_file
from transformers import (
    AutoModelForImageTextToText,
    AutoTokenizer,
    PreTrainedModel,
    Sam3Model,
    Sam3Processor,
)

from .configuration_mlrs import MLRSConfig


logger = logging.getLogger(__name__)

_HUB_KWARGS = ("cache_dir", "force_download", "local_files_only", "proxies", "revision", "token", "subfolder")
_BASE_HUB_KWARGS = ("cache_dir", "force_download", "local_files_only", "proxies", "token")
_NON_MODEL_INPUT_KEYS = ("sam_inputs", "sam_original_sizes", "image_original_sizes")


# ---------------------------------------------------------------------------
# Tool-call parsing
# ---------------------------------------------------------------------------

def _is_number_like(v) -> bool:
    try:
        return isfinite(float(v))
    except Exception:
        return False


def _is_valid_box(b) -> bool:
    return isinstance(b, (list, tuple)) and len(b) == 4 and all(_is_number_like(x) for x in b)


def _infer_grid_from_values(values) -> Optional[int]:
    if not values:
        return 999
    vmax = 0.0
    for v in values:
        try:
            fv = float(v)
            if isfinite(fv):
                vmax = max(vmax, abs(fv))
        except Exception:
            continue
    if vmax <= 1.2:
        return 1
    if vmax <= 1000.5:
        return 999
    return None


def _scale_coord_to_pixel(val, grid, size) -> int:
    if grid is None:
        vv = round(float(val))
    elif grid == 1:
        vv = round(float(val) * (size - 1))
    else:
        vv = round(float(val) * (size - 1) / 999.0)
    return max(0, min(size - 1, int(vv)))


def _boxes_to_pixels(raw_bboxes, image_size_wh, grid) -> List[List[int]]:
    w, h = image_size_wh
    out = []
    for b in raw_bboxes:
        if not _is_valid_box(b):
            continue
        x1 = _scale_coord_to_pixel(b[0], grid, w)
        y1 = _scale_coord_to_pixel(b[1], grid, h)
        x2 = _scale_coord_to_pixel(b[2], grid, w)
        y2 = _scale_coord_to_pixel(b[3], grid, h)
        out.append([min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2)])
    return out


def _normalize_label(label) -> str:
    if label is None:
        return ""
    if isinstance(label, str):
        return label.strip()
    if isinstance(label, (list, tuple)):
        return ", ".join(s for s in (str(x).strip() for x in label if x is not None) if s)
    return str(label).strip()


def _segmentation_to_answer_tags(output_text: str) -> str:
    if output_text is None:
        return ""
    m = re.fullmatch(
        r"^\s*(?:<think>\s*(.*?)\s*</think>\s*)?<segmentation>\s*(.*?)\s*</segmentation>\s*$",
        output_text,
        flags=re.DOTALL,
    )
    if not m:
        return output_text
    think = m.group(1) or ""
    body = (m.group(2) or "").strip()
    if body == "":
        return output_text
    if think.strip():
        return f"<think>{think}</think><answer>{body}</answer>"
    return f"<answer>{body}</answer>"


def _segmentation_body_to_dicts(raw: Optional[str]) -> List[Dict[str, Any]]:
    """Parse one dict, a JSON array of dicts, or comma-joined dicts."""
    if raw is None:
        return []
    raw = raw.strip()
    if raw == "":
        return []

    def _strip_quotes_in_box_array(match):
        return match.group(1) + match.group(2).replace('"', "") + match.group(3)

    def _try_load(s):
        try:
            return json.loads(s)
        except json.JSONDecodeError:
            fixed = re.sub(r'("objects"\s*:\s*\[)(.*?)(\])', _strip_quotes_in_box_array, s, flags=re.DOTALL)
            try:
                return json.loads(fixed)
            except json.JSONDecodeError:
                return None

    data = _try_load(raw)
    if data is None:
        data = _try_load("[" + raw + "]")
    if isinstance(data, dict):
        return [data]
    if isinstance(data, list):
        return [d for d in data if isinstance(d, dict)]
    return []


def _dict_to_raw_bboxes(d: Dict[str, Any]) -> List[List[float]]:
    objects = d.get("objects", [])
    if not isinstance(objects, list):
        return []
    raw_bboxes = []
    for obj in objects:
        if _is_valid_box(obj):
            raw_bboxes.append(list(obj))
        elif isinstance(obj, dict) and _is_valid_box(obj.get("bbox")):
            raw_bboxes.append(list(obj["bbox"]))
    return raw_bboxes


def parse_segmentation_call(output_text: str, image_size_wh: Tuple[int, int]) -> List[Tuple[str, List[List[int]]]]:
    """
    Parse `[<think>..</think>]<segmentation>{...}</segmentation>` into
    `[(noun_phrase, [[x1, y1, x2, y2], ...]), ...]` with pixel boxes.
    Returns [] when the output is not a (valid) segmentation call.
    """
    converted = _segmentation_to_answer_tags(output_text)
    m = re.fullmatch(
        r"^\s*(?:<think>\s*(.*?)\s*</think>\s*)?<answer>\s*(.*?)\s*</answer>\s*$",
        converted,
        re.DOTALL,
    )
    if not m:
        return []

    dicts = _segmentation_body_to_dicts((m.group(2) or "").strip())
    if not dicts:
        return []

    # One coordinate grid for all phrases so scaling is consistent.
    per_dict_boxes = [_dict_to_raw_bboxes(d) for d in dicts]
    all_vals = [v for boxes in per_dict_boxes for b in boxes for v in b[:4]]
    grid = _infer_grid_from_values(all_vals)

    phrases = []
    for d, raw_bboxes in zip(dicts, per_dict_boxes):
        label = _normalize_label(d.get("noun phrase", d.get("noun_phrase", "")))
        if not label.strip() or not raw_bboxes:
            continue
        phrase_grid = grid if grid is not None else _infer_grid_from_values([v for b in raw_bboxes for v in b[:4]])
        boxes = _boxes_to_pixels(raw_bboxes, image_size_wh, phrase_grid)
        if boxes:
            phrases.append((label, boxes))
    return phrases


# ---------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------

@dataclass
class MLRSOutput:
    """
    text:    generated text per sample.
    masks:   per sample, bool tensor [K, H, W] (one mask per noun phrase) at
             `sam_original_sizes` resolution. K == 0 for non-segmentation answers.
    phrases: per sample, the K noun phrases sent to SAM3.
    boxes:   per sample, per phrase, the pixel boxes [x1, y1, x2, y2] sent to SAM3.
    """

    text: List[str]
    masks: List[torch.Tensor]
    phrases: List[List[str]] = field(default_factory=list)
    boxes: List[List[List[List[int]]]] = field(default_factory=list)


def _param_device_dtype(module: torch.nn.Module) -> Tuple[torch.device, torch.dtype]:
    device = torch.device("cpu")
    dtype = torch.float32
    for i, p in enumerate(module.parameters()):
        if i == 0:
            device = p.device
        if p.is_floating_point():
            return device, p.dtype
    return device, dtype


def _to_device(batch: Dict[str, Any], device: torch.device, dtype: torch.dtype) -> Dict[str, Any]:
    out = {}
    for k, v in batch.items():
        if torch.is_tensor(v):
            out[k] = v.to(device=device, dtype=dtype) if v.is_floating_point() else v.to(device)
        else:
            out[k] = v
    return out


class MLRS(PreTrainedModel):
    config_class = MLRSConfig
    base_model_prefix = "mlrs"

    def __init__(
        self,
        config: MLRSConfig,
        language_model: Optional[torch.nn.Module] = None,
        sam_model: Optional[torch.nn.Module] = None,
        tokenizer=None,
        sam_processor=None,
    ):
        super().__init__(config)
        if language_model is None or sam_model is None or tokenizer is None or sam_processor is None:
            raise ValueError("MLRS must be created with MLRS.from_pretrained(...)")
        self.language_model = language_model
        self.sam_model = sam_model
        self.tokenizer = tokenizer
        self.sam_processor = sam_processor

    # -- loading ------------------------------------------------------------

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path,
        *model_args,
        config: Optional[MLRSConfig] = None,
        dtype=None,
        torch_dtype=None,
        device_map=None,
        merge_lora: bool = False,
        **kwargs,
    ) -> "MLRS":
        """
        Loads the InternVL3.5 and SAM3 base models at the revisions pinned in the
        config, attaches the LoRA adapter and overwrites the SAM3 mask decoder.

        merge_lora: merge the adapter into the base weights (faster generation,
            tiny numerical differences vs. the reported results).
        """
        hub_kwargs = {k: kwargs.pop(k) for k in _HUB_KWARGS if k in kwargs}
        kwargs.pop("trust_remote_code", None)
        unused = [k for k in kwargs if not k.startswith("_")]
        if unused:
            logger.warning("MLRS.from_pretrained ignores arguments: %s", unused)
        base_hub_kwargs = {k: hub_kwargs[k] for k in _BASE_HUB_KWARGS if k in hub_kwargs}

        dtype = dtype if dtype is not None else torch_dtype
        if dtype is None or dtype == "auto":
            dtype = torch.bfloat16
        elif isinstance(dtype, str):
            dtype = getattr(torch, dtype)

        if config is None:
            config = MLRSConfig.from_pretrained(pretrained_model_name_or_path, **hub_kwargs)

        path = str(pretrained_model_name_or_path)
        if os.path.isdir(path):
            local_dir = path
        else:
            local_dir = snapshot_download(
                repo_id=path,
                allow_patterns=[f"{config.language_adapter_subdir}/*", config.sam_delta_file, "*.json"],
                **{k: v for k, v in hub_kwargs.items() if k != "subfolder"},
            )

        # Language side: base InternVL3.5 + LoRA adapter.
        # The adapter is kept in its saved precision (autocast_adapter_dtype=True,
        # PEFT's default), matching the setup used for the reported results.
        language_base = AutoModelForImageTextToText.from_pretrained(
            config.language_base_model,
            revision=config.language_base_revision,
            dtype=dtype,
            device_map=device_map,
            **base_hub_kwargs,
        )
        language_model = PeftModel.from_pretrained(
            language_base,
            os.path.join(local_dir, config.language_adapter_subdir),
            is_trainable=False,
        )
        if merge_lora:
            language_model = language_model.merge_and_unload()

        # SAM side: base SAM3 + finetuned weights (mask decoder).
        sam_model = Sam3Model.from_pretrained(
            config.sam_base_model,
            revision=config.sam_base_revision,
            dtype=dtype,
            device_map=device_map,
            **base_hub_kwargs,
        )
        sam_delta = load_file(os.path.join(local_dir, config.sam_delta_file))
        incompatible = sam_model.load_state_dict(sam_delta, strict=False)
        if incompatible.unexpected_keys:
            raise RuntimeError(f"SAM3 delta has keys not present in {config.sam_base_model}: {incompatible.unexpected_keys}")

        tokenizer = AutoTokenizer.from_pretrained(
            config.language_base_model,
            revision=config.language_base_revision,
            **base_hub_kwargs,
        )
        sam_processor = Sam3Processor.from_pretrained(
            config.sam_base_model,
            revision=config.sam_base_revision,
            **base_hub_kwargs,
        )

        model = cls(
            config,
            language_model=language_model,
            sam_model=sam_model,
            tokenizer=tokenizer,
            sam_processor=sam_processor,
        )
        return model.eval()

    def get_input_embeddings(self):
        return self.language_model.get_input_embeddings()

    def forward(self, *args, **kwargs):
        """Language-model forward pass (logits). Use `generate` for text + masks."""
        return self.language_model(*args, **kwargs)

    # -- inference ------------------------------------------------------------

    def _special_token_ids(self, generation_kwargs: Dict[str, Any]) -> Tuple[Optional[int], Optional[int]]:
        eos = generation_kwargs.get("eos_token_id")
        pad = generation_kwargs.get("pad_token_id")
        for cfg in (getattr(self.language_model, "generation_config", None), getattr(self.language_model, "config", None)):
            if cfg is None:
                continue
            if eos is None:
                eos = getattr(cfg, "eos_token_id", None)
            if pad is None:
                pad = getattr(cfg, "pad_token_id", None)
        if isinstance(eos, (list, tuple)):
            eos = eos[0] if eos else None
        return eos, pad

    @staticmethod
    def _generated_mask(tokens: torch.Tensor, eos: Optional[int], pad: Optional[int]) -> torch.Tensor:
        """Valid generated tokens: up to and including the first EOS."""
        mask = torch.zeros_like(tokens, dtype=torch.bool)
        for i, row in enumerate(tokens):
            if eos is not None:
                eos_pos = (row == eos).nonzero(as_tuple=False)
                if eos_pos.numel() > 0:
                    mask[i, : int(eos_pos[0].item()) + 1] = True
                    continue
            mask[i] = row != pad if pad is not None else True
        return mask

    def _is_valid_sam_text(self, text: str) -> bool:
        if not text.strip():
            return False
        max_len = self.sam_model.config.text_config.max_position_embeddings
        n_tokens = self.sam_processor.tokenizer(
            text, truncation=False, add_special_tokens=True, return_tensors="pt"
        )["input_ids"].shape[1]
        return int(n_tokens) <= int(max_len)

    def _segment(
        self,
        sam_inputs: Dict[str, Any],
        row: int,
        size_hw: Tuple[int, int],
        phrases: List[Tuple[str, List[List[int]]]],
    ) -> torch.Tensor:
        """Run SAM3 once per noun phrase of sample `row`. Returns bool [K, H, W]."""
        device, dtype = _param_device_dtype(self.sam_model)
        h, w = size_hw
        k = len(phrases)
        labels = [p[0] for p in phrases]
        boxes = [p[1] for p in phrases]

        base = {}
        for key, v in sam_inputs.items():
            if torch.is_tensor(v):
                v = v.index_select(0, torch.tensor([row], device=v.device))
                v = v.expand(k, *v.shape[1:]).contiguous()
            base[key] = v

        prompt_inputs = self.sam_processor(
            original_sizes=torch.tensor([[h, w]] * k, dtype=torch.long),
            text=labels,
            input_boxes=boxes,
            input_boxes_labels=[[1] * len(bb) for bb in boxes],
            return_tensors="pt",
        )
        base.update(_to_device(dict(prompt_inputs), device, dtype))

        logits = self.sam_model(**base).semantic_seg
        if logits is None or logits.ndim != 4 or logits.shape[1] != 1:
            raise RuntimeError("SAM3 did not return semantic_seg logits of shape [K, 1, Hm, Wm]")
        logits = F.interpolate(logits.float(), size=(h, w), mode="bilinear", align_corners=False)

        if self.config.bbox_filtering:
            inside = torch.zeros_like(logits)
            for i, phrase_boxes in enumerate(boxes):
                for x1, y1, x2, y2 in phrase_boxes:
                    inside[i, :, y1:y2, x1:x2] = 1
            logits = logits * inside - 1000 * (1 - inside)

        return (torch.sigmoid(logits) > self.config.mask_threshold).reshape(k, h, w)

    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.Tensor,
        sam_inputs: Dict[str, Any],
        sam_original_sizes: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        image_original_sizes: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> MLRSOutput:
        """
        Args:
            input_ids, attention_mask, pixel_values, sam_inputs, sam_original_sizes:
                the output of `MLRSProcessor`.
            **kwargs: other tensors are forwarded to the language model; everything
                else is treated as a generation argument (max_new_tokens, do_sample, ...).
        """
        lm_inputs: Dict[str, Any] = {"input_ids": input_ids}
        if attention_mask is not None:
            lm_inputs["attention_mask"] = attention_mask
        generation_kwargs: Dict[str, Any] = {}
        for key, value in kwargs.items():
            if key in _NON_MODEL_INPUT_KEYS:
                continue
            (lm_inputs if torch.is_tensor(value) else generation_kwargs)[key] = value

        generation_kwargs.setdefault("max_new_tokens", self.config.max_new_tokens)
        generation_kwargs.setdefault("do_sample", False)
        generation_kwargs.setdefault("repetition_penalty", 1.0)
        generation_kwargs["num_return_sequences"] = 1

        # 1) Language model.
        lm_device, lm_dtype = _param_device_dtype(self.language_model)
        lm_inputs = _to_device(lm_inputs, lm_device, lm_dtype)
        sequences = self.language_model.generate(**lm_inputs, **generation_kwargs)
        if not torch.is_tensor(sequences):
            sequences = sequences.sequences

        generated = sequences[:, lm_inputs["input_ids"].shape[1]:]
        eos, pad = self._special_token_ids(generation_kwargs)
        valid = self._generated_mask(generated, eos, pad)
        texts = [
            self.tokenizer.decode(ids[m], skip_special_tokens=True)
            for ids, m in zip(generated, valid)
        ]

        # 2) Segmentation tool.
        sam_device, sam_dtype = _param_device_dtype(self.sam_model)
        sam_inputs = _to_device(dict(sam_inputs), sam_device, sam_dtype)
        sizes = sam_original_sizes.to(torch.long).tolist()

        all_masks, all_phrases, all_boxes = [], [], []
        for row, text in enumerate(texts):
            h, w = sizes[row]
            phrases = [
                (label.lstrip(), boxes)
                for label, boxes in parse_segmentation_call(text, image_size_wh=(w, h))
                if self._is_valid_sam_text(label.lstrip())
            ]
            if phrases:
                masks = self._segment(sam_inputs, row, (h, w), phrases)
            else:
                masks = torch.zeros((0, h, w), dtype=torch.bool, device=sam_device)
            all_masks.append(masks)
            all_phrases.append([p[0] for p in phrases])
            all_boxes.append([p[1] for p in phrases])

        return MLRSOutput(text=texts, masks=all_masks, phrases=all_phrases, boxes=all_boxes)