from pathlib import Path
from typing import Dict, List, Tuple

import cv2
import pandas as pd
import torch
from torch.utils.data import Dataset


def normalize_label(value: object) -> float:
    if isinstance(value, str):
        normalized = value.strip().lower()
        if normalized in {"misogyny", "misogynous", "1", "true"}:
            return 1.0
        if normalized in {"not-misogyny", "not misogyny", "non-misogyny", "0", "false"}:
            return 0.0
        raise ValueError(f"Unsupported label value: {value!r}")
    return float(value)


class MemeDataset(Dataset):
    def __init__(
        self,
        split: str,
        data_root: Path,
        language: str,
        tokenizer,
        image_id_col: str,
        text_col: str,
        label_col: str,
        max_length: int,
        oversample_positive_train: int = 0,
    ) -> None:
        self.split = split
        self.split_path = data_root / language / split
        self.tokenizer = tokenizer

        csv_path = self.split_path / f"{split}.csv"
        if not csv_path.exists():
            raise FileNotFoundError(f"Missing CSV file: {csv_path}")

        df = pd.read_csv(csv_path).dropna(axis=0).copy()
        required_cols = {image_id_col, text_col, label_col}
        missing_cols = required_cols.difference(df.columns)
        if missing_cols:
            raise ValueError(f"{csv_path} is missing columns: {sorted(missing_cols)}")

        df[image_id_col] = df[image_id_col].astype(int)
        df[label_col] = df[label_col].map(normalize_label)
        id_to_row = {int(row[image_id_col]): row for _, row in df.iterrows()}

        raw_images: List[torch.Tensor] = []
        raw_records: List[Tuple[int, str, float]] = []
        for image_file in sorted(self.split_path.iterdir()):
            if image_file.suffix.lower() == ".csv":
                continue
            if image_file.suffix.lower() not in {".jpg", ".jpeg", ".png", ".bmp", ".webp"}:
                continue

            try:
                image_id = int(image_file.stem)
            except ValueError:
                continue

            if image_id not in id_to_row:
                continue

            image = cv2.imread(str(image_file))
            if image is None:
                raise ValueError(f"Could not read image: {image_file}")

            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            image = cv2.resize(image, (224, 224), interpolation=cv2.INTER_CUBIC)
            image_tensor = torch.tensor(image, dtype=torch.float32).permute(2, 0, 1) / 255.0
            raw_images.append(image_tensor)

            row = id_to_row[image_id]
            raw_records.append((image_id, str(row[text_col]), float(row[label_col])))

        if not raw_records:
            raise ValueError(f"No valid examples found in {self.split_path}")

        image_stack = torch.stack(raw_images, dim=0)
        mean = image_stack.mean(dim=(0, 2, 3), keepdim=True)
        std = image_stack.std(dim=(0, 2, 3), keepdim=True).clamp_min(1e-6)
        image_stack = (image_stack - mean) / std

        self.images: List[torch.Tensor] = []
        self.labels: List[float] = []
        self.image_ids: List[int] = []
        texts: List[str] = []

        for idx, (image_id, text, label) in enumerate(raw_records):
            repeat = 1
            if split == "train" and int(label) == 1 and oversample_positive_train > 0:
                repeat += oversample_positive_train

            for _ in range(repeat):
                self.images.append(image_stack[idx])
                self.labels.append(label)
                self.image_ids.append(image_id)
                texts.append(text)

        tokens = tokenizer(
            texts,
            truncation=True,
            max_length=max_length,
            padding="max_length",
            return_tensors="pt",
        )
        self.input_ids = tokens.input_ids
        self.attention_mask = tokens.attention_mask

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

    def __getitem__(self, index: int) -> Dict[str, torch.Tensor]:
        return {
            "image": self.images[index],
            "label": torch.tensor(self.labels[index], dtype=torch.float32),
            "input_ids": self.input_ids[index],
            "attention_mask": self.attention_mask[index],
            "image_id": torch.tensor(self.image_ids[index], dtype=torch.long),
        }