from __future__ import annotations

from typing import Any, Protocol
import torch
from torch.utils.data import Dataset

from conveyer.config import ConveyerConfig


class TokenizerProtocol(Protocol):
    vocab_size: int
    pad_id: int
    bos_id: int
    eos_id: int

    def encode(self, text: str) -> list[int]: ...
    def decode(self, ids: list[int]) -> str: ...
    def id_to_piece(self, token_id: int) -> str: ...


def build_dataset_and_tokenizer(
    cfg: ConveyerConfig,
    infer_only: bool = False,
) -> tuple[Dataset[dict[str, Any]], Dataset[dict[str, Any]], TokenizerProtocol]:
    if cfg.data.dataset == "synthetic":
        from conveyer.data.synthetic import build_synthetic_datasets

        return build_synthetic_datasets(cfg)
    elif cfg.data.dataset == "tinystories":
        from conveyer.data.tinystories import build_tinystories_datasets

        return build_tinystories_datasets(cfg, infer_only=infer_only)
    else:
        raise ValueError(f"Unknown dataset '{cfg.data.dataset}'. Choose 'synthetic' or 'tinystories'.")


def collate_conveyer_batch(batch: list[dict[str, Any]], pad_id: int = 0) -> dict[str, torch.Tensor]:
    max_len = max(item["input_ids"].size(0) for item in batch)
    bsz = len(batch)
    input_ids = torch.full((bsz, max_len), pad_id, dtype=torch.long)
    prompt_lengths = torch.zeros(bsz, dtype=torch.long)
    task_ids = torch.zeros(bsz, dtype=torch.long)

    for i, item in enumerate(batch):
        seq = item["input_ids"]
        input_ids[i, : seq.size(0)] = seq
        prompt_lengths[i] = int(item.get("prompt_len", 1))
        task_ids[i] = int(item.get("task_id", 0))

    return {
        "input_ids": input_ids,
        "prompt_lengths": prompt_lengths,
        "task_ids": task_ids,
    }