import json
import random
from typing import List, Dict, Tuple, Iterator
import torch
from torch.utils.data import Dataset, Sampler
from transformers import PreTrainedTokenizerBase


class ConceptSynonymDataset(Dataset):
    """
    Dataset for SapBERT metric learning.
    Each item is a (synonym_text, label_idx) pair.
    The concept_id is mapped to a contiguous integer label [0, N-1].
    """

    def __init__(self, jsonl_path: str):
        self.samples: List[Tuple[str, int]] = []  # (synonym_text, class_label)
        self.concept_to_label: Dict[str, int] = {}
        self.label_to_concept: Dict[int, str] = {}
        self.label_to_indices: Dict[int, List[int]] = {}
        self.concept_names: Dict[str, str] = {}
        self.vocabularies: Dict[str, str] = {}
        self.concept_codes: Dict[str, str] = {}

        self._load_data(jsonl_path)

    def _load_data(self, jsonl_path: str):
        label_counter = 0
        with open(jsonl_path, "r", encoding="utf-8") as f:
            for line in f:
                if not line.strip():
                    continue
                record = json.loads(line)
                cid = str(record["concept_id"])
                synonyms = record.get("synonyms", [])
                if len(synonyms) < 2:
                    continue

                if cid not in self.concept_to_label:
                    label = label_counter
                    self.concept_to_label[cid] = label
                    self.label_to_concept[label] = cid
                    self.concept_names[cid] = record.get("concept_name", "")
                    self.vocabularies[cid] = record.get("vocabulary_id", "")
                    self.concept_codes[cid] = record.get("concept_code", "")
                    label_counter += 1
                else:
                    label = self.concept_to_label[cid]

                if label not in self.label_to_indices:
                    self.label_to_indices[label] = []

                for syn in synonyms:
                    if syn and len(syn.strip()) > 0:
                        idx = len(self.samples)
                        self.samples.append((syn.strip(), label))
                        self.label_to_indices[label].append(idx)

        print(
            f"Loaded {len(self.samples):,} terms across {len(self.concept_to_label):,} unique concepts from {jsonl_path}"
        )

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

    def __getitem__(self, idx: int) -> Tuple[str, int]:
        return self.samples[idx]

    @property
    def num_classes(self) -> int:
        return len(self.concept_to_label)


class MPerClassSampler(Sampler[int]):
    """
    Samples M items per class in each batch.
    Essential for metric learning (e.g., Multi-Similarity Loss) so that every
    concept in a batch has at least M positive pairs.
    """

    def __init__(
        self,
        dataset: ConceptSynonymDataset,
        m: int = 2,
        batch_size: int = 64,
        length_before_new_iter: int = 100000,
    ):
        super().__init__()
        self.dataset = dataset
        self.m = m
        self.batch_size = batch_size
        self.classes_per_batch = max(1, batch_size // m)
        self.labels = list(self.dataset.label_to_indices.keys())
        self.length = min(length_before_new_iter, len(self.dataset))
        self.total_batches = self.length // self.batch_size

    def __len__(self) -> int:
        return self.total_batches * self.batch_size

    def __iter__(self) -> Iterator[int]:
        batches_yielded = 0
        while batches_yielded < self.total_batches:
            selected_classes = random.sample(self.labels, self.classes_per_batch)
            batch = []
            for cls in selected_classes:
                indices = self.dataset.label_to_indices[cls]
                if len(indices) >= self.m:
                    batch.extend(random.sample(indices, self.m))
                else:
                    batch.extend(random.choices(indices, k=self.m))

            random.shuffle(batch)
            for idx in batch:
                yield idx
            batches_yielded += 1


def collate_fn(batch: List[Tuple[str, int]], tokenizer: PreTrainedTokenizerBase, max_length: int = 40):
    """
    Collate function to dynamically tokenize batch texts and convert labels to tensors.
    """
    texts = [item[0] for item in batch]
    labels = torch.tensor([item[1] for item in batch], dtype=torch.long)

    encoded = tokenizer(
        texts,
        padding=True,
        truncation=True,
        max_length=max_length,
        return_tensors="pt",
    )

    return {
        "input_ids": encoded["input_ids"],
        "attention_mask": encoded["attention_mask"],
        "labels": labels,
    }