import json
from typing import List
from torch.utils.data import Dataset
import logging
from jinja2 import Template
import re

from dataset.data_utils import build_messages

logger = logging.getLogger(__name__)


def extract_prompt(full_prompt: str, pattern) -> str | None:
    match = re.search(pattern, full_prompt, re.DOTALL)
    if match:
        return match.group(1).strip()
    return None


class TongkuanListWiseCotSFTParallelDataset(Dataset):
    def __init__(self, args, data_path: str, train=True) -> None:
        super().__init__()
        self.full_crop_mode = getattr(args, "full_crop_mode", False)
        if isinstance(data_path, str):
            if data_path.endswith(".json"):
                self.data = json.load(open(data_path))
            elif data_path.endswith(".jsonl"):
                self.data = [json.loads(line) for line in open(data_path)]
            else:
                raise ValueError(f"Unsupported data file: {data_path}")
        elif isinstance(data_path, List):
            self.data = data_path
        else:
            raise ValueError(f"Unsupported data_path type: {type(data_path)}")

        print(f"Load {len(self.data)} instances")
        self.args = args
        self.train = train
        with open(args.format_prompt, encoding="utf-8") as f:
            format_prompt = f.read()
        format_prompt = Template(format_prompt.strip())
        full_prompt = format_prompt.render()
        self.sys_prompt = extract_prompt(full_prompt, r"<sys_start>(.*?)<sys_end>")
        self.user_prompt = extract_prompt(full_prompt, r"<user_start>(.*?)<user_end>")
        self.max_full_pixels = args.max_full_tokens * (args.patch_size * 2) ** 2
        self.max_crop_pixels = args.max_crop_tokens * (args.patch_size * 2) ** 2

    def __len__(self) -> int:
        return len(self.data)

    def _get_docs(self, item):
        docs = list(item["docs"])
        if self.train:
            return docs[:self.args.num_train_docs]
        return docs

    def get_instance(self, idx):
        item = self.data[idx]
        query_image = item["query_image"]
        label_mapping = {0: 0, 1: 1, 2: 2, 3: 3}
        docs_to_process = self._get_docs(item)
        final_doc_labels = [label_mapping.get(doc.get("label", None)) for doc in docs_to_process]

        if self.full_crop_mode:
            query = [
                {"image": query_image, "max_pixels": self.max_full_pixels},
                {"image": item["query_crop_image"], "max_pixels": self.max_crop_pixels},
            ]
            docs = []
            for doc in docs_to_process:
                docs.append({"image": doc["doc_image"], "max_pixels": self.max_full_pixels})
                docs.append({"image": doc["crop_image"], "max_pixels": self.max_crop_pixels})
        else:
            query = [{"image": query_image, "max_pixels": self.max_crop_pixels}]
            docs = [{"image": doc["crop_image"], "max_pixels": self.max_crop_pixels} for doc in docs_to_process]

        return {
            "query_image": query,
            "doc_images": docs,
            "labels": final_doc_labels,
            "titles": [doc["doc_title"] for doc in docs_to_process],
            "instruction": item.get("instruction", ""),
        }

    def _build_cot_target(self, item, labels):
        cot = item.get("cot")
        answer = item.get("cot_answer") or "".join(str(label) for label in labels)
        if not cot:
            raise ValueError("COT training data must contain a non-empty `cot` field.")
        if len(answer) != len(labels) or any(ch not in "0123" for ch in answer):
            raise ValueError(f"Invalid cot_answer: {answer}")
        cot = cot.strip()
        if "【最终答案】" in cot or "最终答案" in cot:
            return cot
        return f"{cot}\n\n【最终答案】\n{answer}"

    def __getitem__(self, i):
        item = self.data[i]
        instance = self.get_instance(i)
        query_image = instance["query_image"]
        doc_images = instance["doc_images"]
        labels = instance["labels"]
        item["labels"] = labels
        query_image = [query_image] if not isinstance(query_image, List) else query_image
        images = query_image + doc_images

        raw_res = self._build_cot_target(item, labels) if self.train else None
        rerank_message = build_messages(
            self.sys_prompt,
            self.user_prompt,
            images,
            raw_res,
            titles=list(instance["titles"]),
            instruction=instance["instruction"],
        )
        return rerank_message, item