import json
from typing import Dict, List
from torch.utils.data import Dataset
import random
import logging
from jinja2 import Template
import re
from .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:
        sys_prompt = match.group(1).strip()
        return sys_prompt
    else:
        return None


class TongkuanListWiseSFTParallelDataset(Dataset):

    def __init__(
        self,
        args,
        data_path: str,
        train=True
    ) -> None:
        super(TongkuanListWiseSFTParallelDataset, self).__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)]
        if isinstance(data_path, List):
            self.data = 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
        self.class_mode = getattr(args, "class_mode", 0)

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

    def _get_train_docs(self, item, num_train_docs):
        if 'docs' in item:
            docs_to_process = list(item['docs'])
            random.shuffle(docs_to_process)
            docs_to_process = docs_to_process[:num_train_docs]
        return docs_to_process

    def get_instance(self, idx):
        item = self.data[idx]

        query_image = item['query_image']

        label_mapping = {0: 0, 1: 1, 2: 2, 3: 3}

        if self.train:
            docs_to_process = self._get_train_docs(item, self.args.num_train_docs)
        else:
            docs_to_process = item['docs']

        final_doc_labels = []

        for doc in docs_to_process:
            final_doc_labels.append(label_mapping.get(doc.get('label', None)))

        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]

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

        return instance

    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
        if self.train:
            raw_res = "".join(str(l) for l in labels)
            rerank_message = build_messages(self.sys_prompt, self.user_prompt, images, raw_res,
                titles=list(instance['titles']), instruction=instance['instruction'])
        else:
            rerank_message = build_messages(self.sys_prompt, self.user_prompt, images, None,
                titles=list(instance['titles']), instruction=instance['instruction'])
        return rerank_message, item