use std::time::Instant;

/// Records a single generation run and computes TPS.
#[derive(Debug)]
pub struct GenerationRecord {
    pub start: Instant,
    pub prompt_tokens: u32,
    pub completion_tokens: u32,
}

impl GenerationRecord {
    pub fn new(prompt_tokens: u32) -> Self {
        Self {
            start: Instant::now(),
            prompt_tokens,
            completion_tokens: 0,
        }
    }

    pub fn tick(&mut self) {
        self.completion_tokens += 1;
    }

    pub fn elapsed_secs(&self) -> f64 {
        self.start.elapsed().as_secs_f64()
    }

    pub fn tokens_per_second(&self) -> f64 {
        if self.completion_tokens == 0 {
            return 0.0;
        }
        self.completion_tokens as f64 / self.elapsed_secs()
    }
}