"""
Supervised learning dataset implementations from HuggingFace datasets.
"""

import json
from typing import Any, Callable

import blobfile
import chz
import datasets
import tinker
import torch
from tinker_cookbook.renderers import Message, Renderer, TrainOnWhat
from tinker_cookbook.supervised.common import datum_from_model_input_weights
from tinker_cookbook.supervised.types import ChatDatasetBuilder, SupervisedDataset


def conversation_to_datum(
    conversation: list[Message],
    renderer: Renderer,
    max_length: int | None,
    train_on_what: TrainOnWhat = TrainOnWhat.ALL_ASSISTANT_MESSAGES,
) -> tinker.Datum:
    """Common function to process a list of messages into a Datum."""
    model_input, weights = renderer.build_supervised_example(
        conversation, train_on_what=train_on_what
    )
    return datum_from_model_input_weights(model_input, weights, max_length)


def text_to_datum(text: str, renderer: Renderer, max_length: int | None) -> tinker.Datum:
    # This just trains on all tokens (train_on_what is ignored for simple text)
    tokens = list(renderer.tokenizer.encode(text, add_special_tokens=False))
    weights = torch.ones(len(tokens), dtype=torch.float32)
    # Convert tokens to ModelInput with a single text chunk
    model_input = tinker.ModelInput(chunks=[tinker.types.EncodedTextChunk(tokens=tokens)])
    return datum_from_model_input_weights(model_input, weights, max_length)


def _one_of(a: Any, b: Any) -> bool:
    return (a is not None and b is None) or (a is None and b is not None)


class SupervisedDatasetFromHFDataset(SupervisedDataset):
    def __init__(
        self,
        hf_dataset: datasets.Dataset,
        batch_size: int,
        map_fn: Callable[[dict], tinker.Datum] | None = None,
        flatmap_fn: Callable[[dict], list[tinker.Datum]] | None = None,
    ):
        assert _one_of(map_fn, flatmap_fn), "Only one of map_fn or flatmap_fn can be provided"
        self.hf_dataset = hf_dataset
        self.shuffle_dataset = (
            hf_dataset  # Keep a reference to the original dataset to avoid statefulness
        )
        self.batch_size = batch_size
        self.map_fn = map_fn
        self.flatmap_fn = flatmap_fn

    def get_batch(self, index: int) -> list[tinker.Datum]:
        rows = self.shuffle_dataset.select(
            range(index * self.batch_size, (index + 1) * self.batch_size)
        )
        if self.map_fn is not None:
            return [self.map_fn(row) for row in rows.to_list()]
        else:
            assert self.flatmap_fn is not None
            return [datum for row in rows.to_list() for datum in self.flatmap_fn(row)]

    def set_epoch(self, seed: int = 0):
        self.shuffle_dataset = self.hf_dataset.shuffle(seed=seed)

    def __len__(self) -> int:
        return len(self.hf_dataset) // self.batch_size


class StreamingSupervisedDatasetFromHFDataset(SupervisedDataset):
    def __init__(
        self,
        hf_dataset: datasets.IterableDataset,
        batch_size: int,
        length: int,
        map_fn: Callable[[dict], tinker.Datum] | None = None,
        flatmap_fn: Callable[[dict], list[tinker.Datum]] | None = None,
        buffer_size: int = 10_000,
    ):
        assert _one_of(map_fn, flatmap_fn), "Only one of map_fn or flatmap_fn can be provided"
        self.hf_dataset = hf_dataset.shuffle(seed=0, buffer_size=buffer_size).batch(
            batch_size=batch_size, drop_last_batch=True
        )
        self.dataset_iterator = iter(self.hf_dataset)
        self.index = -1
        self.batch_size = batch_size
        self.map_fn = map_fn
        self.flatmap_fn = flatmap_fn
        # We pass the length to the dataset, since streaming HF datasets don't have a length attribute
        self.length = length

    def get_batch(self, index: int) -> list[tinker.Datum]:
        # TODO: this is a hack to make sure the index is correct
        # should maybe think about a more robust way to do this
        assert index == self.index + 1
        self.index = index
        batch = next(self.dataset_iterator)
        rows = [dict(zip(batch.keys(), values)) for values in zip(*batch.values())]
        if self.map_fn is not None:
            return [self.map_fn(row) for row in rows]
        else:
            assert self.flatmap_fn is not None
            return [datum for row in rows for datum in self.flatmap_fn(row)]

    def set_epoch(self, seed: int = 0):
        self.hf_dataset.set_epoch(seed)
        self.dataset_iterator = iter(self.hf_dataset)
        self.index = -1

    def __len__(self) -> int:
        return self.length // self.batch_size


@chz.chz
class FromConversationFileBuilder(ChatDatasetBuilder):
    file_path: str
    limit: int | None = None
    test_size: int = 0
    shuffle_seed: int | None = 42  # None means no shuffle

    def __call__(self) -> tuple[SupervisedDataset, SupervisedDataset | None]:
        # Load conversations from JSONL file
        conversations = []
        with blobfile.BlobFile(self.file_path, "r", streaming=False) as f:
            for line in f:
                data = json.loads(line.strip())
                if "messages" not in data:
                    raise ValueError(
                        f"Each line in the JSONL file must contain a 'messages' field. Got: {data.keys()}"
                    )
                conversations.append(data)
                if self.limit is not None and len(conversations) >= self.limit:
                    break

        # Create HuggingFace dataset from the loaded data
        dataset = datasets.Dataset.from_list(conversations)

        # Shuffle if seed is provided
        if self.shuffle_seed is not None:
            dataset = dataset.shuffle(seed=self.shuffle_seed)

        # Split into train and test
        if self.test_size > 0 and len(dataset) > self.test_size:
            test_ds = dataset.take(self.test_size)
            train_ds = dataset.skip(self.test_size)
        else:
            # If test_size is 0 or dataset is too small, use all data for training
            train_ds = dataset
            test_ds = None

        # Use train_on_what from common_config if provided, otherwise use default
        train_on_what = (
            TrainOnWhat(self.common_config.train_on_what)
            if self.common_config.train_on_what
            else TrainOnWhat.ALL_ASSISTANT_MESSAGES
        )

        # Define mapping function
        def map_fn(row: dict) -> tinker.Datum:
            return conversation_to_datum(
                row["messages"], self.renderer, self.common_config.max_length, train_on_what
            )

        # Create supervised dataset
        supervised_dataset = SupervisedDatasetFromHFDataset(
            train_ds, batch_size=self.common_config.batch_size, map_fn=map_fn
        )

        # Create evaluator if we have test data
        if test_ds is not None:
            test_dataset = SupervisedDatasetFromHFDataset(
                test_ds, batch_size=len(test_ds), map_fn=map_fn
            )
        else:
            test_dataset = None

        return supervised_dataset, test_dataset


def _load_text_or_messages_file(
    file_path: str, limit: int | None = None, shuffle_seed: int | None = None
) -> datasets.Dataset:
    """Load a JSONL file containing 'messages' or 'text' fields into a HuggingFace Dataset."""
    conversations = []
    with blobfile.BlobFile(file_path, "r", streaming=False) as f:
        for line in f:
            data = json.loads(line.strip())
            if "messages" not in data and "text" not in data:
                raise ValueError(
                    f"Each line in the JSONL file must contain a 'messages' or 'text' field. Got: {data.keys()}"
                )
            # Normalize the data structure - ensure both fields exist to prevent HF Dataset from dropping columns
            # HF Dataset infers schema and will set missing columns to None
            normalized_data = {}
            if "messages" in data and data["messages"] is not None:
                normalized_data["messages"] = data["messages"]
                normalized_data["text"] = None  # Explicitly set to None
            elif "text" in data and data["text"] is not None:
                normalized_data["text"] = data["text"]
                normalized_data["messages"] = None  # Explicitly set to None
            else:
                # Skip rows where both are None or missing
                continue

            conversations.append(normalized_data)
            if limit is not None and len(conversations) >= limit:
                break

    # Create HuggingFace dataset from the loaded data
    dataset = datasets.Dataset.from_list(conversations)
    print(f"Loaded {len(dataset)} rows from {file_path}")

    # Shuffle if seed is provided
    if shuffle_seed is not None:
        print(f"Shuffling with seed: {shuffle_seed}")
        dataset = dataset.shuffle(seed=shuffle_seed)
    else:
        print("NOTE: Not shuffling since shuffle_seed is None")

    return dataset


@chz.chz
class FromTextOrMessagesFileBuilder(FromConversationFileBuilder):
    file_path: str
    limit: int | None = None
    test_file_path: str | None = None
    shuffle_seed: int | None = 0  # None means no shuffle

    def __call__(self) -> tuple[SupervisedDataset, SupervisedDataset | None]:
        # Load train dataset
        train_ds = _load_text_or_messages_file(
            self.file_path, limit=self.limit, shuffle_seed=self.shuffle_seed
        )

        # Load test dataset if path provided
        if self.test_file_path is not None:
            test_ds = _load_text_or_messages_file(
                self.test_file_path, limit=self.limit, shuffle_seed=self.shuffle_seed
            )
        else:
            test_ds = None

        # Use train_on_what from common_config if provided, otherwise use default
        train_on_what = (
            TrainOnWhat(self.common_config.train_on_what)
            if self.common_config.train_on_what
            else TrainOnWhat.ALL_ASSISTANT_MESSAGES
        )

        # Define mapping function
        def map_fn(row: dict) -> tinker.Datum:
            # HuggingFace datasets returns a special dict-like object
            # Check for messages first
            if "messages" in row:
                messages = row["messages"]
                # Check if actually not None (HF might set None for missing columns)
                if messages is not None:
                    return conversation_to_datum(
                        messages, self.renderer, self.common_config.max_length, train_on_what
                    )

            # Then check for text
            if "text" in row:
                text = row["text"]
                # Check if actually not None (HF might set None for missing columns)
                if text is not None:
                    assert isinstance(text, str), f"Text must be a string. Got: {type(text)}"
                    return text_to_datum(text, self.renderer, self.common_config.max_length)

            raise ValueError(
                f"Row must contain either 'messages' or 'text' with non-None values. Got: {row.keys()}"
            )

        # Create supervised dataset
        supervised_dataset = SupervisedDatasetFromHFDataset(
            train_ds, batch_size=self.common_config.batch_size, map_fn=map_fn
        )

        # Create evaluator if we have test data
        if test_ds is not None:
            test_dataset = SupervisedDatasetFromHFDataset(
                test_ds, batch_size=len(test_ds), map_fn=map_fn
            )
        else:
            test_dataset = None

        return supervised_dataset, test_dataset
