"""
Hugging Face Dataset Ingestion & Instruction/Chat Formatting Pipeline for Nova 1.0
Supports raw text datasets (e.g. Wikipedia) and instruction/conversational datasets (e.g. Alpaca, Dolly).
"""

import os
from typing import Iterator, List, Tuple
import torch
from torch.utils.data import Dataset, DataLoader
from datasets import load_dataset
from src.tokenizer.bpe_tokenizer import Nova1Tokenizer


class Nova1TokenizedDataset(Dataset):
    def __init__(self, token_chunks: List[List[int]], pad_token_id: int = 0):
        self.chunks = token_chunks
        self.pad_token_id = pad_token_id

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

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
        chunk = self.chunks[idx]
        # Language modeling task: input = chunk[:-1], target = chunk[1:]
        input_ids = torch.tensor(chunk[:-1], dtype=torch.long)
        target_ids = torch.tensor(chunk[1:], dtype=torch.long)
        return input_ids, target_ids


def format_item_to_chat(item: dict) -> str:
    """Format diverse HF dataset schemas into standardized User/Assistant dialogue format."""
    # Format 1: Dolly OAI messages format [{'role': 'user', 'content': ...}, {'role': 'assistant', 'content': ...}]
    if "messages" in item and isinstance(item["messages"], list):
        formatted = ""
        for msg in item["messages"]:
            role = "User" if msg.get("role") == "user" else "Assistant"
            content = msg.get("content", "").strip()
            formatted += f"{role}: {content}\n"
        return formatted.strip()

    # Format 2: Alpaca format {'instruction': ..., 'input': ..., 'output': ...}
    if "instruction" in item and "output" in item:
        instruction = item["instruction"].strip()
        inp = item.get("input", "").strip()
        output = item["output"].strip()
        if inp:
            user_text = f"{instruction}\nContext: {inp}"
        else:
            user_text = instruction
        return f"User: {user_text}\nAssistant: {output}"

    # Format 3: Standard raw text field
    text_val = item.get("text") or item.get("content") or list(item.values())[0]
    return str(text_val)


def create_hf_dataloader(
    dataset_name: str,
    tokenizer: Nova1Tokenizer,
    max_seq_len: int = 256,
    batch_size: int = 16,
    split: str = "train",
    subset: str = None,
    max_samples: int = 10000,
    shuffle: bool = True
) -> DataLoader:
    """
    Downloads/loads Hugging Face dataset, formats chat/instruction dialogue, tokenizes text,
    packs into chunks, and returns PyTorch DataLoader.
    """
    token = os.environ.get("HF_TOKEN")
    if subset:
        raw_ds = load_dataset(dataset_name, subset, split=split, token=token)
    else:
        raw_ds = load_dataset(dataset_name, split=split, token=token)

    if max_samples and len(raw_ds) > max_samples:
        raw_ds = raw_ds.select(range(max_samples))

    # Prepend high-frequency core conversational dialogues so Nova 1.0 masters greetings
    core_dialogues = [
        "User: Hello!\nAssistant: Hello! How can I help you today?",
        "User: Hi!\nAssistant: Hi there! How can I assist you today?",
        "User: How are you?\nAssistant: I am doing well, thank you! How can I help you?",
        "User: What is your name?\nAssistant: I am Nova 1.0, an AI assistant built from scratch.",
        "User: Hello, who are you?\nAssistant: Hello! I am Nova 1.0, a custom AI model created using the Hierarchical Reasoning Model architecture."
    ]

    all_tokens = []
    chunk_size = max_seq_len + 1  # Need max_seq_len + 1 for input/target split

    # Repeat core dialogues to ensure high weight in training
    for _ in range(500):
        for dlg in core_dialogues:
            ids = tokenizer.encode(dlg, add_bos=True, add_eos=True)
            all_tokens.extend(ids)

    for item in raw_ds:
        formatted_text = format_item_to_chat(item)
        if not formatted_text or len(formatted_text.strip()) == 0:
            continue
        ids = tokenizer.encode(formatted_text, add_bos=True, add_eos=True)
        all_tokens.extend(ids)

    # Chunk tokens into max_seq_len blocks
    token_chunks = []
    for i in range(0, len(all_tokens) - chunk_size + 1, chunk_size):
        chunk = all_tokens[i : i + chunk_size]
        token_chunks.append(chunk)

    if len(token_chunks) == 0:
        raise ValueError(f"Dataset produced 0 chunks for max_seq_len={max_seq_len}. Increase dataset sample size.")

    dataset = Nova1TokenizedDataset(token_chunks, pad_token_id=tokenizer.pad_id)
    return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=True)