import json
import torch
from torch.utils.data import Dataset
from typing import List, Dict, Tuple
from nlp_core.tokenizer import BilingualTokenizer
from nlp_core.vocabulary import Vocabulary

class SummarizationDataset(Dataset):
    """
    PyTorch Dataset for Sequence-to-Sequence Summarization.
    """

    def __init__(
        self,
        data_path: str,
        src_vocab: Vocabulary,
        trg_vocab: Vocabulary,
        lang: str = "en",
        max_src_len: int = 120,
        max_trg_len: int = 40
    ):
        self.src_vocab = src_vocab
        self.trg_vocab = trg_vocab
        self.lang = lang
        self.max_src_len = max_src_len
        self.max_trg_len = max_trg_len
        self.tokenizer = BilingualTokenizer()

        with open(data_path, "r", encoding="utf-8") as f:
            self.samples = json.load(f)

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

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
        item = self.samples[idx]
        article = item["article"]
        summary = item["summary"]

        # Tokenize
        src_tokens = self.tokenizer.tokenize_words(article, lang=self.lang, remove_stopwords=False)
        trg_tokens = self.tokenizer.tokenize_words(summary, lang=self.lang, remove_stopwords=False)

        # Encode to indices
        src_indices = self.src_vocab.encode(src_tokens, add_sos=True, add_eos=True, max_len=self.max_src_len)
        trg_indices = self.trg_vocab.encode(trg_tokens, add_sos=True, add_eos=True, max_len=self.max_trg_len)

        return torch.tensor(src_indices, dtype=torch.long), torch.tensor(trg_indices, dtype=torch.long)


def pad_collate_fn(batch: List[Tuple[torch.Tensor, torch.Tensor]], pad_idx: int = 0) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Collates a list of (src, trg) tensor pairs and dynamically pads them
    to the max length within the mini-batch.
    """
    src_list, trg_list = zip(*batch)

    # Pad source tensors
    src_padded = torch.nn.utils.rnn.pad_sequence(src_list, batch_first=True, padding_value=pad_idx)

    # Pad target tensors
    trg_padded = torch.nn.utils.rnn.pad_sequence(trg_list, batch_first=True, padding_value=pad_idx)

    return src_padded, trg_padded