import torch
from torch.utils.data import Dataset
from tokenizers import Tokenizer
from typing import List, Dict

from src.common.prompts import format_prompt, resolve_eos_id, resolve_pad_id

class SFTDataset(Dataset):
    """
    Dataset for Supervised Fine-Tuning.
    Formats Alpaca/Instruction data by masking the prompt loss and only computing loss on the response.
    """
    def __init__(self, data: List[Dict[str, str]], tokenizer: Tokenizer, max_len: int = 512):
        self.samples = []

        pad_token_id = resolve_pad_id(tokenizer)
        eos_token_id = resolve_eos_id(tokenizer)

        for item in data:
            # Canonical shared prompt format (keeps SFT/DPO/inference aligned).
            prompt = format_prompt(item["prompt"])
            response = item["response"]

            enc_prompt = tokenizer.encode(prompt).ids
            enc_response = tokenizer.encode(response).ids
            if eos_token_id is not None:
                enc_response = enc_response + [eos_token_id]
            
            input_ids = enc_prompt + enc_response
            # -100 tells PyTorch's CrossEntropyLoss to ignore these tokens during loss calculation
            labels = [-100] * len(enc_prompt) + enc_response
            
            # Truncate or pad
            if len(input_ids) > max_len:
                input_ids = input_ids[:max_len]
                labels = labels[:max_len]
            else:
                pad_len = max_len - len(input_ids)
                input_ids = input_ids + [pad_token_id] * pad_len
                labels = labels + [-100] * pad_len
                
            self.samples.append((
                torch.tensor(input_ids, dtype=torch.long),
                torch.tensor(labels, dtype=torch.long)
            ))

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

    def __getitem__(self, idx):
        return self.samples[idx]