import json
import re

import torch
from typing_extensions import override

import comfy.model_management
import comfy.sd
import comfy.supported_models_base
import comfy.utils
import folder_paths
from comfy.text_encoders.hunyuan_video import llama_detect
from comfy_api.latest import ComfyExtension, io

from .wevisone_model import DECODER_MEMORY, SYSTEM_PROMPT, WeVisOneTEModel, WeVisOneTokenizer

TASKS = ["detection", "referring", "point", "count_box", "count_point", "matting", "ocr", "ocr_grounding", "freeform"]
REF_RE = re.compile(r"<ref>(.*?)</ref>\s*<(box|point|polygon)>(.*?)(?:</(?:box|point|polygon)>|<ref>|$)", re.DOTALL)
COORD_RE = re.compile(r"<(\d{1,4})>")


def build_prompt(task, query):
    labels = [line.strip() for line in query.splitlines() if line.strip()]
    text = " ".join(labels)
    phrase = text if not text or text[-1] in ".?!" else text + "."
    if task == "detection":
        return "Locate all the instances that match the following description: " + ". ".join(labels) + "."
    if task == "referring":
        return "Locate the instances that match the following description: " + phrase
    if task == "point":
        return "point to: " + ". ".join(labels)
    if task == "count_box":
        return "How many {} in the image? Please count them with boxes.".format(text or "objects")
    if task == "count_point":
        return "How many {} in the image?".format(text or "objects")
    if task == "matting":
        return "Locate the instances that match the following description and provide the alpha mattes: " + phrase
    if task == "ocr":
        return "detect all the word in the box format"
    if task == "ocr_grounding":
        return "Locate the following text in the image: " + text
    return query


def parse_response(text, width, height):
    boxes, points, seg_boxes = [], [], []
    for label, shape, body in REF_RE.findall(text):
        label = label.strip() or "object"
        for segment in body.split(","):
            values = [int(v) for v in COORD_RE.findall(segment)]
            if shape == "point":
                for i in range(0, len(values) - 1, 2):
                    points.append({"x": round(values[i] / 1000 * width), "y": round(values[i + 1] / 1000 * height), "label": label})
            elif shape == "box":
                for i in range(0, len(values) - 3, 4):
                    x1, y1, x2, y2 = values[i:i + 4]
                    if x2 > x1 and y2 > y1:
                        boxes.append({"x": x1 / 1000 * width, "y": y1 / 1000 * height, "width": (x2 - x1) / 1000 * width, "height": (y2 - y1) / 1000 * height, "label": label})
    # Mask queries: the last four coordinates before each <OBJ> inside a box.
    for body in re.findall(r"<box>(.*?)(?:</box>|$)", text, re.DOTALL):
        for prefix in body.split("<OBJ>")[:-1]:
            values = [int(v) for v in COORD_RE.findall(prefix)]
            if len(values) >= 4:
                x1, y1, x2, y2 = values[-4:]
                if x2 > x1 and y2 > y1 and max(x1, y1, x2, y2) <= 1000:
                    seg_boxes.append([x1 / 1000, y1 / 1000, x2 / 1000, y2 / 1000])
    return boxes, points, seg_boxes


class WeVisOneVLLoader(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="WeVisOneVLLoader",
            display_name="WeVisOne-VL Loader",
            category="WeVisOne-VL",
            inputs=[io.Combo.Input("clip_name", options=folder_paths.get_filename_list("text_encoders"))],
            outputs=[io.Clip.Output()],
        )

    @classmethod
    def execute(cls, clip_name) -> io.NodeOutput:
        sd = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("text_encoders", clip_name), safe_load=True)
        target = comfy.supported_models_base.ClipTarget(WeVisOneTokenizer, WeVisOneTEModel)
        detect = llama_detect(sd)
        model_options = {"dtype": detect["dtype_llama"]}
        if "llama_quantization_metadata" in detect:
            model_options["quantization_metadata"] = detect["llama_quantization_metadata"]
        clip = comfy.sd.CLIP(target, embedding_directory=folder_paths.get_folder_paths("embeddings"), parameters=comfy.utils.calculate_parameters(sd),
                             state_dict=[sd], model_options=model_options)
        return io.NodeOutput(clip)


class WeVisOneVLPerceive(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="WeVisOneVLPerceive",
            display_name="WeVisOne-VL Perceive",
            category="WeVisOne-VL",
            search_aliases=["detect", "grounding", "segment", "matting", "ocr", "count", "point"],
            inputs=[
                io.Clip.Input("clip"),
                io.Image.Input("image"),
                io.Combo.Input("task", options=TASKS, default="matting"),
                io.String.Input("query", multiline=True, default="", tooltip="Object labels (one per line for detection/point), a referring phrase, text to find, or the full prompt for freeform."),
                io.Int.Input("max_length", default=1024, min=1, max=8192),
                io.Boolean.Input("individual_masks", default=False, tooltip="Output one alpha matte per instance instead of their union."),
            ],
            outputs=[
                io.String.Output(display_name="text"),
                io.BoundingBox.Output(display_name="bboxes"),
                io.String.Output(display_name="points", tooltip="JSON [{\"x\": int, \"y\": int}, ...] in pixels, as SAM3 point prompts."),
                io.Mask.Output(display_name="mask", tooltip="Alpha mattes for the matting task, or filled boxes for the other box tasks: their union per image, or every instance with individual_masks."),
            ],
        )

    @classmethod
    def execute(cls, clip, image, task, query, max_length, individual_masks=False) -> io.NodeOutput:
        height, width = image.shape[1:3]
        prompt = build_prompt(task, query)
        texts, all_boxes, all_points, masks = [], [], [], []
        for i in range(image.shape[0]):
            frame = image[i:i + 1]
            tokens = clip.tokenize(prompt, image=frame, system_prompt=SYSTEM_PROMPT)
            # top_k=1 is greedy decoding that still applies the reference repetition penalty (ignored by do_sample=False).
            ids = clip.generate(tokens, do_sample=True, max_length=max_length, temperature=1.0, top_k=1, top_p=1.0, min_p=0.0, repetition_penalty=1.05, seed=0)
            text = clip.decode(ids, skip_special_tokens=False).split("<|im_end|>", 1)[0].split("<|endoftext|>", 1)[0].strip()
            boxes, points, seg_boxes = parse_response(text, width, height)
            mattes = torch.zeros((0, height, width))
            if seg_boxes:
                comfy.model_management.load_models_gpu([clip.patcher], memory_required=DECODER_MEMORY)
                clip.cond_stage_model.set_clip_options({"execution_device": clip.patcher.load_device})
                mattes = clip.cond_stage_model.decode_masks(tokens, ids, frame, seg_boxes).cpu()
            elif boxes:
                mattes = torch.zeros((len(boxes), height, width))
                for m, b in zip(mattes, boxes):
                    m[round(b["y"]):round(b["y"] + b["height"]), round(b["x"]):round(b["x"] + b["width"])] = 1.0
            texts.append(text)
            all_boxes.append(boxes)
            all_points.extend({"x": p["x"], "y": p["y"]} for p in points)
            masks.append(mattes if individual_masks else mattes.amax(0, keepdim=True) if len(mattes) else torch.zeros((1, height, width)))
        return io.NodeOutput("\n".join(texts), all_boxes, json.dumps(all_points), torch.cat(masks))


class WeVisOneVLExtension(ComfyExtension):
    @override
    async def get_node_list(self) -> list[type[io.ComfyNode]]:
        return [WeVisOneVLLoader, WeVisOneVLPerceive]


async def comfy_entrypoint() -> WeVisOneVLExtension:
    return WeVisOneVLExtension()