import json
import importlib
import re
import sys
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from transformers import AutoConfig, AutoProcessor, PreTrainedModel, Qwen3VLForConditionalGeneration
from transformers.utils import cached_file

try:
    from .configuration_samtok import SAMTokCompositeConfig
except ImportError:
    from configuration_samtok import SAMTokCompositeConfig


class SAMTokCompositeModel(PreTrainedModel):
    config_class = SAMTokCompositeConfig
    main_input_name = "input_ids"
    _no_split_modules = ["Qwen3VLDecoderLayer", "MultiScaleBlock", "TwoWayAttentionBlock"]

    def __init__(self, config):
        super().__init__(config)
        self.vlm = None
        self.processor = None
        self.vq_sam2 = None
        self.sam2_image_processor = None
        self._samtok_models = None

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        config = kwargs.pop("config", None)
        trust_remote_code = kwargs.pop("trust_remote_code", True)
        dtype = kwargs.pop("dtype", kwargs.pop("torch_dtype", "auto"))
        device_map = kwargs.pop("device_map", None)
        vlm_attn_implementation = kwargs.pop("vlm_attn_implementation", None)

        hub_kwargs = {
            key: kwargs.pop(key)
            for key in list(kwargs.keys())
            if key in {"cache_dir", "force_download", "local_files_only", "revision", "token"}
        }

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

        model = cls(config)
        model._load_samtok_package(pretrained_model_name_or_path)
        qwen_config_path = cached_file(pretrained_model_name_or_path, config.qwen_config_file, **hub_kwargs)
        with open(qwen_config_path, "r", encoding="utf-8") as f:
            qwen_config_dict = json.load(f)
        qwen_config = AutoConfig.for_model(qwen_config_dict.pop("model_type"), **qwen_config_dict)

        vlm_kwargs = dict(kwargs)
        vlm_kwargs["config"] = qwen_config
        vlm_kwargs["dtype"] = dtype
        if device_map is not None:
            vlm_kwargs["device_map"] = device_map
        if vlm_attn_implementation is not None:
            vlm_kwargs["attn_implementation"] = vlm_attn_implementation

        model.vlm = Qwen3VLForConditionalGeneration.from_pretrained(
            pretrained_model_name_or_path,
            *model_args,
            **vlm_kwargs,
        ).eval()
        model.processor = AutoProcessor.from_pretrained(
            pretrained_model_name_or_path,
            trust_remote_code=trust_remote_code,
            **hub_kwargs,
        )
        model.vq_sam2 = model._load_vq_sam2(pretrained_model_name_or_path, hub_kwargs).eval()
        return model

    def _load_samtok_package(self, pretrained_model_name_or_path):
        model_path = Path(str(pretrained_model_name_or_path))
        if model_path.exists():
            model_dir = model_path.resolve()
            if str(model_dir) not in sys.path:
                sys.path.insert(0, str(model_dir))
        self._samtok_models = importlib.import_module("samtok.models")
        self.sam2_image_processor = self._samtok_models.DirectResize(self.config.sam2_image_size)

    @property
    def device(self):
        if self.vlm is not None:
            return self.vlm.device
        return super().device

    @property
    def dtype(self):
        if self.vlm is not None:
            return self.vlm.dtype
        return super().dtype

    def to(self, *args, **kwargs):
        super().to(*args, **kwargs)
        if self.vlm is not None:
            self.vlm.to(*args, **kwargs)
        if self.vq_sam2 is not None:
            self.vq_sam2.to(*args, **kwargs)
        return self

    def eval(self):
        super().eval()
        if self.vlm is not None:
            self.vlm.eval()
        if self.vq_sam2 is not None:
            self.vq_sam2.eval()
        return self

    def forward(self, *args, **kwargs):
        return self.vlm(*args, **kwargs)

    def generate(self, *args, **kwargs):
        return self.vlm.generate(*args, **kwargs)

    def _load_vq_sam2(self, pretrained_model_name_or_path, hub_kwargs):
        if self._samtok_models is None:
            self._load_samtok_package(pretrained_model_name_or_path)
        sam2_ckpt_path = cached_file(
            pretrained_model_name_or_path,
            self.config.sam2_ckpt_file,
            **hub_kwargs,
        )
        vq_sam2_path = cached_file(
            pretrained_model_name_or_path,
            self.config.vq_sam2_file,
            **hub_kwargs,
        )
        sam2_config = self._samtok_models.SAM2Config(ckpt_path=sam2_ckpt_path)
        vq_sam2_config = self._samtok_models.VQ_SAM2Config(
            sam2_config=sam2_config,
            codebook_size=self.config.codebook_size,
            codebook_depth=self.config.codebook_depth,
            shared_codebook=self.config.shared_codebook,
            latent_dim=self.config.latent_dim,
        )
        vq_sam2 = self._samtok_models.VQ_SAM2(vq_sam2_config)
        state = torch.load(vq_sam2_path, map_location="cpu", weights_only=False)
        vq_sam2.load_state_dict(state)
        if self.vlm is not None:
            vq_sam2.to(self.vlm.device)
        return vq_sam2

    def generate_with_masks(
        self,
        image,
        question,
        max_new_tokens=512,
        do_sample=False,
        top_p=1.0,
        **generate_kwargs,
    ):
        image_obj = _load_image(image)
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image},
                    {"type": "text", "text": question},
                ],
            }
        ]
        inputs = self.processor.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_dict=True,
            return_tensors="pt",
        )
        inputs = inputs.to(self.device)
        generated_ids = self.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=do_sample,
            top_p=top_p,
            **generate_kwargs,
        )
        generated_ids_trimmed = [
            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
        ]
        text = self.processor.batch_decode(
            generated_ids_trimmed,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=False,
        )[0]
        masks, tags = self.decode_masks(text, image_obj)
        return {"text": text, "masks": masks, "tags": tags}

    def decode_masks(self, output_text, image):
        image_obj = _load_image(image)
        ori_width, ori_height = image_obj.size
        quant_ids = _extract_mt_token_ids_v1(output_text)
        if len(quant_ids) % self.config.codebook_depth != 0:
            output_text = _fix_mt_format_comprehensive(output_text)
            quant_ids = _extract_mt_token_ids_v2(output_text)
        if len(quant_ids) == 0 or len(quant_ids) % self.config.codebook_depth != 0:
            return None, []

        remap_quant_ids = []
        tags = []
        for idx in range(len(quant_ids) // self.config.codebook_depth):
            chunk = quant_ids[
                idx * self.config.codebook_depth:(idx + 1) * self.config.codebook_depth
            ]
            normalized = [
                token_id - book_id * self.config.codebook_size
                for book_id, token_id in enumerate(chunk)
            ]
            normalized = [
                code if 0 <= code < self.config.codebook_size else -1
                for code in normalized
            ]
            if all(code == -1 for code in normalized):
                continue
            tags.append("-".join(str(x) for x in chunk))
            remap_quant_ids.append(normalized)

        if not remap_quant_ids:
            return None, []
        remap_quant_ids = remap_quant_ids[: self.config.max_samtok_masks]
        tags = tags[: self.config.max_samtok_masks]

        sam2_image = self.sam2_image_processor.apply_image(np.array(image_obj))
        pixel_values = torch.from_numpy(sam2_image).permute(2, 0, 1).contiguous()
        pixel_values = pixel_values.unsqueeze(0).to(self.vq_sam2.dtype).to(self.vq_sam2.device)
        quant_ids_tensor = torch.LongTensor(remap_quant_ids).to(self.vq_sam2.device)
        pred_masks = self._forward_masks_in_chunks(pixel_values, quant_ids_tensor)
        pred_masks = torch.nn.functional.interpolate(
            pred_masks,
            size=(ori_height, ori_width),
            mode="bilinear",
        )
        pred_masks = pred_masks > 0.5
        pred_masks = pred_masks[:, 0, :, :].cpu().numpy().astype(np.uint8)
        return pred_masks, tags

    def _forward_masks_in_chunks(self, pixel_values, quant_ids):
        chunks = []
        chunk_size = self.config.decode_chunk_size
        for start in range(0, len(quant_ids), chunk_size):
            quant_chunk = quant_ids[start:start + chunk_size]
            pixel_chunk = pixel_values.repeat(len(quant_chunk), 1, 1, 1)
            with torch.no_grad():
                chunks.append(self.vq_sam2.forward_with_codes(pixel_chunk, quant_chunk))
        return torch.cat(chunks, dim=0)


def _load_image(image):
    if isinstance(image, Image.Image):
        return image.convert("RGB")
    return Image.open(Path(image)).convert("RGB")


def _extract_mt_token_ids_v1(text):
    return [int(x) for x in re.findall(r"<\|mt_(\d{4})\|>", text)]


def _extract_mt_token_ids_v2(text):
    pattern = re.compile(r"<\|mt_start\|><\|mt_(\d{4})\|><\|mt_(\d{4})\|><\|mt_end\|>")
    values = []
    for first, second in pattern.findall(text):
        values.extend([int(first), int(second)])
    return values


def _fix_mt_format_comprehensive(text):
    text = re.sub(
        r"(<\|mt_start\|>)(<\|mt_\d+\|>)(<\|mt_\d+\|>)(?:<\|mt_\d+\|>)+<\|mt_end\|>",
        r"\1\2\3<|mt_end|>",
        text,
    )
    text = re.sub(
        r"(<\|mt_start\|>)(<\|mt_\d+\|>)(<\|mt_end\|>)",
        r"\1\2<|mt_9999|><|mt_end|>",
        text,
    )
    text = re.sub(
        r"(<\|mt_start\|>)(<\|mt_\d+\|>)(?!<\|mt_)",
        r"\1\2<|mt_9999|><|mt_end|>",
        text,
    )
    return text