//! Data pipeline: load the GPT-2 BPE (reusing xserv's from-scratch tokenizer),
//! tokenize a text corpus into one flat token stream, and sample fixed-length
//! `(input, target)` windows for next-token prediction. Host-only (no GPU).

use std::path::Path;
use xserv_tokenizer::Tokenizer;

/// A tokenized corpus: one flat stream of token ids, plus the vocab size.
pub struct Corpus {
    pub tokens: Vec<i32>,
    pub vocab_size: usize,
}

impl Corpus {
    /// Load `tokenizer.json` (GPT-2 BPE) and tokenize the UTF-8 text at
    /// `corpus_path` into a single id stream. TinyStories separates stories with
    /// `<|endoftext|>`; the GPT-2 tokenizer emits that as a single special token,
    /// so document boundaries are preserved in the stream.
    pub fn load(tokenizer_path: &Path, corpus_path: &Path) -> Self {
        let tok = Tokenizer::from_file(tokenizer_path);
        let text = std::fs::read_to_string(corpus_path)
            .unwrap_or_else(|e| panic!("failed to read corpus {}: {e}", corpus_path.display()));
        // The range-fetched corpus may start/end mid-story; drop a leading partial
        // line and a trailing partial story so we only train on whole sentences.
        let text = trim_to_whole_stories(&text);
        let ids: Vec<i32> = tok.encode(text).into_iter().map(|t| t as i32).collect();
        Self {
            tokens: ids,
            vocab_size: tok.vocab_size(),
        }
    }

    /// Total number of tokens.
    pub fn len(&self) -> usize {
        self.tokens.len()
    }

    pub fn is_empty(&self) -> bool {
        self.tokens.is_empty()
    }

    /// Sample one `(input, target)` pair of length `seq` for next-token
    /// prediction: a window `[s, s+seq+1)` → input `[s, s+seq)`, target shifted
    /// by one. `rng_state` is advanced (a tiny LCG, so sampling is reproducible
    /// from a seed without pulling in an RNG crate).
    pub fn sample(&self, seq: usize, rng_state: &mut u64) -> (Vec<i32>, Vec<i32>) {
        assert!(self.tokens.len() > seq + 1, "corpus shorter than a window");
        let max_start = self.tokens.len() - seq - 1;
        let start = (next_rand(rng_state) % (max_start as u64 + 1)) as usize;
        let input = self.tokens[start..start + seq].to_vec();
        let target = self.tokens[start + 1..start + seq + 1].to_vec();
        (input, target)
    }
}

/// Drop a leading partial line (before the first newline) and everything after
/// the last `<|endoftext|>` marker, so a byte-range download still yields only
/// complete stories. Falls back to the raw text if no marker is present.
fn trim_to_whole_stories(text: &str) -> &str {
    let start = text.find('\n').map(|i| i + 1).unwrap_or(0);
    let body = &text[start..];
    match body.rfind("<|endoftext|>") {
        Some(end) => &body[..end + "<|endoftext|>".len()],
        None => body,
    }
}

/// Tiny LCG (same constants as the model tests' deterministic fill) so dataset
/// sampling is reproducible from a single u64 seed.
fn next_rand(state: &mut u64) -> u64 {
    *state = state
        .wrapping_mul(6364136223846793005)
        .wrapping_add(1442695040888963407);
    *state >> 16
}