"""Dataset and collator for Delta causal language modeling."""

from __future__ import annotations

import json
import logging
import os
from pathlib import Path
from typing import Any

import torch
from torch.utils.data import Dataset

from delta.tokenizer import DeltaTokenizer

logging.basicConfig(level=os.getenv("DELTA_LOG_LEVEL", "INFO").upper())
logger = logging.getLogger(__name__)


def _read_texts(path: Path) -> list[str]:
    """Read .txt files or .jsonl files containing a text field."""

    files = [path] if path.is_file() else sorted(path.rglob("*"))
    texts: list[str] = []
    for file_path in files:
        if file_path.suffix.lower() == ".txt":
            texts.append(file_path.read_text(encoding="utf-8"))
        elif file_path.suffix.lower() == ".jsonl":
            with file_path.open("r", encoding="utf-8") as handle:
                for line in handle:
                    line = line.strip()
                    if not line:
                        continue
                    record = json.loads(line)
                    if "text" in record:
                        texts.append(str(record["text"]))
    return texts


class DeltaDataset(Dataset[dict[str, torch.Tensor]]):
    """Sliding-window token dataset for language modeling."""

    def __init__(
        self,
        data_path: str | Path,
        tokenizer: DeltaTokenizer,
        max_seq_len: int = 512,
        stride: int = 256,
    ) -> None:
        self.data_path = Path(data_path)
        self.tokenizer = tokenizer
        self.max_seq_len = max_seq_len
        self.stride = stride
        texts = _read_texts(self.data_path)
        if not texts:
            raise ValueError(f"No .txt or .jsonl text records found in {self.data_path}")
        self.windows: list[list[int]] = []
        for text in texts:
            ids = tokenizer.encode(text, add_special_tokens=True)
            for start in range(0, max(1, len(ids) - 1), stride):
                window = ids[start : start + max_seq_len]
                if len(window) >= 2:
                    self.windows.append(window)
                if start + max_seq_len >= len(ids):
                    break
        logger.info("Loaded %s training windows from %s", len(self.windows), self.data_path)

    def __len__(self) -> int:
        """Return the number of windows."""

        return len(self.windows)

    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
        """Return one token window."""

        ids = torch.tensor(self.windows[index], dtype=torch.long)
        return {"input_ids": ids, "labels": ids.clone()}


class DeltaDataCollator:
    """Dynamic padding collator for causal language modeling."""

    def __init__(self, pad_token_id: int = 0) -> None:
        self.pad_token_id = pad_token_id

    def __call__(self, features: list[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
        """Pad input ids and labels to the longest sample in the batch."""

        max_len = max(feature["input_ids"].size(0) for feature in features)
        input_ids = torch.full((len(features), max_len), self.pad_token_id, dtype=torch.long)
        labels = torch.full((len(features), max_len), -100, dtype=torch.long)
        attention_mask = torch.zeros((len(features), max_len), dtype=torch.long)
        for row, feature in enumerate(features):
            ids = feature["input_ids"]
            length = ids.size(0)
            input_ids[row, :length] = ids
            labels[row, :length] = feature["labels"]
            pad_positions = ids == self.pad_token_id
            labels[row, :length][pad_positions] = -100
            attention_mask[row, :length] = 1
        return {"input_ids": input_ids, "labels": labels, "attention_mask": attention_mask}