"""Data pipeline for SplitBit LLM training.

Streaming data loader — doesn't load full dataset into RAM.
Dynamic batching based on available memory.
Data augmentation: random cropping, context window sliding.
"""

from __future__ import annotations

import logging
import os
import random
from typing import Iterator, List

import numpy as np

from ..model.tokenizer import BPETokenizer, BOS_ID, EOS_ID

logger = logging.getLogger(__name__)


class DataPipeline:
    """Streaming text data pipeline for training.

    Reads .txt files one at a time, encodes to tokens, and yields
    batches of fixed-length sequences for training.
    """

    def __init__(
        self,
        tokenizer: BPETokenizer,
        seq_len: int = 256,
        batch_size: int = 4,
        data_dir: str | None = None,
    ) -> None:
        self.tokenizer = tokenizer
        self.seq_len = seq_len
        self.batch_size = batch_size
        self.data_dir = data_dir

    def load_texts(self, paths: List[str] | str) -> List[str]:
        """Load text from files. Returns list of text strings."""
        if isinstance(paths, str):
            paths = [paths]

        texts = []
        for path in paths:
            if os.path.isdir(path):
                for fname in sorted(os.listdir(path)):
                    if fname.endswith((".txt", ".md", ".jsonl")):
                        fpath = os.path.join(path, fname)
                        with open(fpath, "r", encoding="utf-8", errors="replace") as f:
                            texts.append(f.read())
            elif os.path.isfile(path):
                with open(path, "r", encoding="utf-8", errors="replace") as f:
                    texts.append(f.read())
            else:
                logger.warning("Data path not found: %s", path)

        total_chars = sum(len(t) for t in texts)
        logger.info("Loaded %d texts, %d total chars", len(texts), total_chars)
        return texts

    def encode_texts(self, texts: List[str]) -> np.ndarray:
        """Encode all texts into a single token array."""
        all_ids: List[int] = []
        for text in texts:
            ids = self.tokenizer.encode(text, add_bos=True, add_eos=True)
            all_ids.extend(ids)
        logger.info("Encoded %d total tokens", len(all_ids))
        return np.array(all_ids, dtype=np.int64)

    def create_batches(self, token_ids: np.ndarray, shuffle: bool = True) -> Iterator[tuple[np.ndarray, np.ndarray]]:
        """Yield (input, target) batches for next-token prediction.

        input: [batch, seq_len] — tokens 0..seq_len-1
        target: [batch, seq_len] — tokens 1..seq_len (shifted by 1)
        """
        n_tokens = len(token_ids)
        n_seqs = n_tokens // (self.seq_len + 1)
        if n_seqs == 0:
            logger.warning("Not enough tokens (%d) for seq_len=%d", n_tokens, self.seq_len)
            return

        # Truncate to full sequences
        usable = n_seqs * (self.seq_len + 1)
        token_ids = token_ids[:usable]

        # Reshape into sequences
        sequences = token_ids.reshape(n_seqs, self.seq_len + 1)

        # Shuffle
        if shuffle:
            np.random.shuffle(sequences)

        # Yield batches
        for i in range(0, n_seqs, self.batch_size):
            batch = sequences[i:i + self.batch_size]
            if len(batch) < self.batch_size:
                break
            inputs = batch[:, :-1]  # [batch, seq_len]
            targets = batch[:, 1:]   # [batch, seq_len]
            yield inputs, targets

    def stream_batches(self, paths: List[str] | str, shuffle: bool = True) -> Iterator[tuple[np.ndarray, np.ndarray]]:
        """Stream batches from files without loading everything at once.

        Reads one file at a time, encodes, and yields batches.
        """
        if isinstance(paths, str):
            paths = [paths]

        for path in paths:
            texts = self.load_texts(path)
            if not texts:
                continue
            token_ids = self.encode_texts(texts)
            yield from self.create_batches(token_ids, shuffle=shuffle)