use anyhow::Result;
use serde::Deserialize;
use std::path::Path;
use tokenizers::Tokenizer;

#[derive(Debug, Deserialize)]
struct DollyRow {
    instruction: String,
    #[serde(default)]
    context: String,
    response: String,
}

/// A tokenized example. `mask[i] = 1.0` marks the positions that count
/// towards the loss: only the answer, never the prompt.
#[derive(Debug, Clone)]
pub struct Example {
    pub ids: Vec<i64>,
    pub mask: Vec<f32>,
}

fn encode(tok: &Tokenizer, text: &str) -> Result<Vec<i64>> {
    Ok(tok
        .encode(text, false)
        .map_err(|e| anyhow::anyhow!("{e}"))?
        .get_ids()
        .iter()
        .map(|&i| i as i64)
        .collect())
}

pub fn load_dolly(
    path: impl AsRef<Path>,
    tok: &Tokenizer,
    max_len: usize,
    limit: usize,
) -> Result<Vec<Example>> {
    let raw = std::fs::read_to_string(path)?;
    let mut out = Vec::new();

    for line in raw.lines() {
        if out.len() >= limit {
            break;
        }
        let row: DollyRow = match serde_json::from_str(line) {
            Ok(r) => r,
            Err(_) => continue,
        };

        // ChatML, the SmolLM2 chat template
        let user = if row.context.trim().is_empty() {
            row.instruction.clone()
        } else {
            format!("{}\n\n{}", row.instruction, row.context)
        };
        let prompt = format!("<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n");
        let answer = format!("{}<|im_end|>", row.response);

        let p = encode(tok, &prompt)?;
        let a = encode(tok, &answer)?;
        if p.len() + a.len() > max_len || a.is_empty() {
            continue;
        }

        let mut ids = p.clone();
        ids.extend_from_slice(&a);
        let mut mask = vec![0.0; p.len()];
        mask.extend(std::iter::repeat(1.0).take(a.len()));

        out.push(Example { ids, mask });
    }
    Ok(out)
}

/// Pads a batch to its longest example. Padding has mask 0, so it adds
/// nothing to the loss.
pub struct Batch {
    pub ids: Vec<i64>,
    pub mask: Vec<f32>,
    pub batch: usize,
    pub seq: usize,
}

pub fn collate(examples: &[Example]) -> Batch {
    let seq = examples.iter().map(|e| e.ids.len()).max().unwrap_or(1);
    let batch = examples.len();
    let mut ids = Vec::with_capacity(batch * seq);
    let mut mask = Vec::with_capacity(batch * seq);

    for e in examples {
        ids.extend_from_slice(&e.ids);
        mask.extend_from_slice(&e.mask);
        for _ in e.ids.len()..seq {
            ids.push(0);
            mask.push(0.0);
        }
    }
    Batch {
        ids,
        mask,
        batch,
        seq,
    }
}