import os
import torch
import json
import numpy as np
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms


class SimpleImageTextDataset(Dataset):

    def __init__(
        self,
        data_source,          # ✅ 可以是 CSV / list / all.jsonl
        tokenizer,
        image_size=224,
        use_modal_embedding=False,
        default_modal_index=0,
        dataset_root=""       # ✅ 给 all.jsonl 用
    ):
        self.tokenizer = tokenizer
        self.use_modal_embedding = use_modal_embedding
        self.default_modal_index = default_modal_index
        self.image_size = image_size
        self.dataset_root = dataset_root

        # ------- 读取样本 -------
        if isinstance(data_source, str):

            # ✅ 1️⃣ 如果是 CSV
            if data_source.endswith(".csv"):
                import pandas as pd
                df = pd.read_csv(data_source)
                self.samples = list(zip(df["image_path"], df["text"]))

            # ✅ 2️⃣ 如果是 all.jsonl
            elif data_source.endswith(".jsonl"):
                self.samples = self._load_from_jsonl(data_source)

            else:
                raise ValueError("Unsupported file type. Use CSV or JSONL.")

        else:
            # ✅ 3️⃣ 直接是 list [(img_path, text), ...]
            self.samples = data_source

        # ------- 图像 transform -------
        self.transform = transforms.Compose([
            transforms.Grayscale(num_output_channels=1),
            transforms.Resize((image_size, image_size)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.5], std=[0.5])
        ])

        print(f"✅ Loaded {len(self.samples)} image-text pairs.")

    def _load_from_jsonl(self, jsonl_path):
        samples = []
        with open(jsonl_path, "r", encoding="utf-8") as f:
            for line in f:
                item = json.loads(line.strip())

                rel_path = item["img_path"]

                rel_path = rel_path.replace("\\", "/")

                if rel_path.startswith("dataset/"):
                    rel_path = rel_path[len("dataset/"):]

                if self.dataset_root:
                    img_path = os.path.join(self.dataset_root, rel_path)
                else:
                    img_path = rel_path

                text = item["text"].replace("\n", " ").strip()
                samples.append((img_path, text))

        return samples

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, index):
        img_path, text_str = self.samples[index]

        # ---- 1. 载入图像 -> [1,224,224] ----
        img = Image.open(img_path).convert("L")
        img = self.transform(img)

        # ---- 2. tokenize 文本 ----
        text = self.tokenizer(
            text_str,
            return_tensors="pt",
            truncation=True,
            padding="max_length",
            max_length=512,
        )

        text = {k: v.squeeze(0) for k, v in text.items()}

        # ---- 3. modal_embedding（可选）----
        if self.use_modal_embedding:
            modal_index = torch.tensor(self.default_modal_index, dtype=torch.long)
            return {
                "image": img,
                "text": text,
                "modal_index": modal_index
            }

        return {
            "image": img,
            "text": text
        }


class TextBatch:
    def __init__(self, d):
        for k, v in d.items():
            setattr(self, k, v)


def simple_collate(batch):
    images = torch.stack([item["image"] for item in batch], dim=0)

    text_dict = {
        "input_ids": torch.stack([item["text"]["input_ids"] for item in batch]),
        "attention_mask": torch.stack([item["text"]["attention_mask"] for item in batch]),
    }

    if "token_type_ids" in batch[0]["text"]:
        text_dict["token_type_ids"] = torch.stack(
            [item["text"]["token_type_ids"] for item in batch]
        )

    text = TextBatch(text_dict)

    if "modal_index" in batch[0]:
        modal_indexs = torch.stack([item["modal_index"] for item in batch])
        return images, text, modal_indexs

    return images, text