"""
PyTorch Datasets and DataLoaders for Document Pre-training and QA Fine-Tuning.
"""
import torch
from torch.utils.data import Dataset
from typing import List, Dict, Any, Optional, Tuple
from tokenizer.tokenizer import CustomBPETokenizer


class DocPretrainDataset(Dataset):
    """
    Dataset for next-token autoregressive training on extracted document text.
    Segments token stream into fixed-length context windows with stride.
    """
    def __init__(
        self,
        text_corpus: List[str],
        tokenizer: CustomBPETokenizer,
        seq_len: int = 256,
        stride: int = 128
    ):
        self.tokenizer = tokenizer
        self.seq_len = seq_len
        self.samples = []

        # Encode full corpus into continuous token list
        full_tokens = []
        for text in text_corpus:
            tokens = tokenizer.encode(text, add_bos=True, add_eos=True)
            full_tokens.extend(tokens)

        # Slice into overlapping sequences
        if len(full_tokens) <= seq_len + 1:
            # Pad if shorter than sequence length
            pad_len = (seq_len + 1) - len(full_tokens)
            sample = full_tokens + [tokenizer.pad_id] * pad_len
            self.samples.append(sample)
        else:
            for i in range(0, len(full_tokens) - seq_len, stride):
                chunk = full_tokens[i : i + seq_len + 1]
                if len(chunk) == seq_len + 1:
                    self.samples.append(chunk)

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

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
        tokens = self.samples[idx]
        x = torch.tensor(tokens[:-1], dtype=torch.long)
        y = torch.tensor(tokens[1:], dtype=torch.long)
        return x, y


class DocQADataset(Dataset):
    """
    Dataset for Context-conditioned Document Question-Answering.
    Formats samples as:
    <bos> <doc> {context} <sep> <query> {question} <sep> <ans> {answer} <eos>
    """
    def __init__(
        self,
        qa_pairs: List[Dict[str, str]],
        tokenizer: CustomBPETokenizer,
        max_seq_len: int = 512,
        mask_prompt_loss: bool = True
    ):
        """
        qa_pairs: List of {"context": str, "question": str, "answer": str}
        mask_prompt_loss: If True, only compute gradient loss on the answer tokens.
        """
        self.tokenizer = tokenizer
        self.max_seq_len = max_seq_len
        self.mask_prompt_loss = mask_prompt_loss
        self.samples = []

        for item in qa_pairs:
            context = item.get("context", "")
            question = item.get("question", "")
            answer = item.get("answer", "")

            # Construct structured formatted string
            prompt_str = f"<bos> <doc> {context} <sep> <query> {question} <sep> <ans> "
            answer_str = f"{answer} <eos>"

            prompt_ids = tokenizer.encode(prompt_str, add_bos=False, add_eos=False)
            answer_ids = tokenizer.encode(answer_str, add_bos=False, add_eos=False)

            total_ids = prompt_ids + answer_ids

            if len(total_ids) > max_seq_len:
                # Truncate context if necessary to fit max sequence length
                max_ctx_len = max_seq_len - len(answer_ids) - 10
                if max_ctx_len > 10:
                    truncated_prompt_str = f"<bos> <doc> {context[:max_ctx_len*3]} <sep> <query> {question} <sep> <ans> "
                    prompt_ids = tokenizer.encode(truncated_prompt_str, add_bos=False, add_eos=False)
                    total_ids = (prompt_ids + answer_ids)[:max_seq_len]
                else:
                    total_ids = total_ids[:max_seq_len]

            self.samples.append({
                "input_ids": total_ids[:-1],
                "target_ids": total_ids[1:],
                "prompt_len": len(prompt_ids) - 1
            })

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

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        item = self.samples[idx]
        x = list(item["input_ids"])
        y = list(item["target_ids"])
        prompt_len = item["prompt_len"]

        # Mask prompt tokens in target labels if required
        if self.mask_prompt_loss:
            # Set target label to -1 (ignored by CrossEntropyLoss) for prompt positions
            for p in range(min(prompt_len, len(y))):
                y[p] = -1

        return {
            "input_ids": torch.tensor(x, dtype=torch.long),
            "target_ids": torch.tensor(y, dtype=torch.long),
            "length": len(x)
        }


def qa_collate_fn(batch: List[Dict[str, Any]], pad_id: int = 0) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Collate variable-length sequences with right-padding.
    """
    max_len = max(item["length"] for item in batch)
    batch_x = []
    batch_y = []

    for item in batch:
        x = item["input_ids"]
        y = item["target_ids"]
        pad_len = max_len - len(x)

        if pad_len > 0:
            padded_x = torch.cat([x, torch.full((pad_len,), pad_id, dtype=torch.long)])
            padded_y = torch.cat([y, torch.full((pad_len,), -1, dtype=torch.long)])
        else:
            padded_x = x
            padded_y = y

        batch_x.append(padded_x)
        batch_y.append(padded_y)

    return torch.stack(batch_x), torch.stack(batch_y)