//! Token sampling and text generation.

use crate::model::transformer::{InferenceSession, LLMHostsModel};

/// Generation parameters.
#[derive(Debug, Clone)]
pub struct GenerateParams {
    pub max_tokens: usize,
    pub temperature: f32,
    pub top_p: f32,
    pub stop_tokens: Vec<u32>,
}

impl Default for GenerateParams {
    fn default() -> Self {
        Self {
            max_tokens: 256,
            temperature: 0.0,
            top_p: 0.9,
            stop_tokens: vec![],
        }
    }
}

/// Sample a token from logits.
///
/// Temperature 0.0 = greedy (argmax).
/// Temperature > 0 + top_p = nucleus sampling.
pub fn sample_token(logits: &[f32], params: &GenerateParams) -> u32 {
    if params.temperature <= 0.0 || params.temperature < 1e-7 {
        // Greedy: argmax
        return logits
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .map(|(i, _)| i as u32)
            .unwrap_or(0);
    }

    // Apply temperature
    let mut probs: Vec<f32> = logits.iter().map(|&l| l / params.temperature).collect();

    // Softmax
    let max_val = probs.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let mut sum = 0.0f32;
    for p in &mut probs {
        *p = (*p - max_val).exp();
        sum += *p;
    }
    for p in &mut probs {
        *p /= sum;
    }

    // Top-p (nucleus) filtering
    let mut indexed: Vec<(usize, f32)> = probs.iter().enumerate().map(|(i, &p)| (i, p)).collect();
    indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

    let mut cumulative = 0.0f32;
    let mut cutoff_idx = indexed.len();
    for (i, (_, p)) in indexed.iter().enumerate() {
        cumulative += p;
        if cumulative >= params.top_p {
            cutoff_idx = i + 1;
            break;
        }
    }
    indexed.truncate(cutoff_idx);

    // Renormalize
    let total: f32 = indexed.iter().map(|(_, p)| p).sum();
    for (_, p) in &mut indexed {
        *p /= total;
    }

    // Sample
    use rand::Rng;
    let mut rng = rand::thread_rng();
    let r: f32 = rng.gen();
    let mut acc = 0.0f32;
    for (idx, prob) in &indexed {
        acc += prob;
        if acc >= r {
            return *idx as u32;
        }
    }

    // Fallback to last
    indexed.last().map(|(i, _)| *i as u32).unwrap_or(0)
}

/// Generate tokens from a loaded model.
///
/// Creates a fresh [`InferenceSession`] (KV cache) for this request and runs
/// the full prefill + decode loop. Takes `model` by shared reference — callers
/// may hold `Arc<LLMHostsModel>` and call this from multiple threads safely.
///
/// `prompt_tokens`: tokenized prompt.
/// Returns generated token IDs (excluding prompt).
pub fn generate(
    model: &LLMHostsModel,
    prompt_tokens: &[u32],
    params: &GenerateParams,
) -> Vec<u32> {
    // Allocate a fresh per-request session (KV cache).
    let mut session = InferenceSession::new(&model.config);

    // Prefill: process all prompt tokens at once
    let logits = model.forward(prompt_tokens, &mut session);
    let vocab_size = model.config.vocab_size;

    // Sample first generated token from last position's logits
    let last_logits = &logits.data[(prompt_tokens.len() - 1) * vocab_size..prompt_tokens.len() * vocab_size];
    let mut next_token = sample_token(last_logits, params);

    let mut generated = Vec::with_capacity(params.max_tokens);

    for _ in 0..params.max_tokens {
        if params.stop_tokens.contains(&next_token) {
            break;
        }
        generated.push(next_token);

        // Decode: process one token at a time
        let logits = model.forward(&[next_token], &mut session);
        let token_logits = &logits.data[..vocab_size];
        next_token = sample_token(token_logits, params);
    }

    generated
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_greedy_sampling_deterministic() {
        let logits = vec![0.0, 1.0, 5.0, 2.0, 3.0];
        let params = GenerateParams { temperature: 0.0, ..Default::default() };
        let token = sample_token(&logits, &params);
        assert_eq!(token, 2); // index of max value (5.0)
    }

    #[test]
    fn test_greedy_sampling_repeated() {
        let logits = vec![1.0, 10.0, 3.0];
        let params = GenerateParams { temperature: 0.0, ..Default::default() };
        // Should always return same result for greedy
        for _ in 0..10 {
            assert_eq!(sample_token(&logits, &params), 1);
        }
    }

    #[test]
    fn test_sampling_with_temperature() {
        let logits = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let params = GenerateParams {
            temperature: 1.0,
            top_p: 1.0,
            ..Default::default()
        };
        // Should return a valid index
        let token = sample_token(&logits, &params);
        assert!((token as usize) < logits.len());
    }

    #[test]
    fn test_top_p_filtering() {
        // Make one logit dominate
        let mut logits = vec![-100.0; 100];
        logits[42] = 100.0;
        let params = GenerateParams {
            temperature: 1.0,
            top_p: 0.5,
            ..Default::default()
        };
        // With top_p=0.5, the dominant token should always be selected
        for _ in 0..10 {
            assert_eq!(sample_token(&logits, &params), 42);
        }
    }

    #[test]
    fn test_sample_empty_logits_returns_zero() {
        let params = GenerateParams::default();
        assert_eq!(sample_token(&[], &params), 0);
    }

    #[test]
    fn test_default_params() {
        let params = GenerateParams::default();
        assert_eq!(params.max_tokens, 256);
        assert_eq!(params.temperature, 0.0);
        assert_eq!(params.top_p, 0.9);
        assert!(params.stop_tokens.is_empty());
    }

    fn load_test_model() -> (crate::model::config::LLMHostsModelConfig, crate::model::transformer::LLMHostsModel) {
        let dir = tempfile::tempdir().unwrap();
        let config_gen = crate::test_utils::TestModelConfig::default();
        let path = crate::test_utils::generate_test_gguf(dir.path(), &config_gen);
        let header = crate::gguf::parse_header(&path).unwrap();
        let mapped = crate::gguf::MappedWeights::open(&path, header.data_offset).unwrap();
        let config = crate::model::config::LLMHostsModelConfig::from_gguf(&header).unwrap();
        let weights = crate::model::transformer::LLMHostsModelWeights::from_gguf(&config, &header, &mapped).unwrap();
        let model = crate::model::transformer::LLMHostsModel::new(config.clone(), weights);
        (config, model)
    }

    #[test]
    fn test_generate_produces_tokens() {
        let (config, model) = load_test_model();
        let params = GenerateParams {
            max_tokens: 10,
            temperature: 0.0,
            ..Default::default()
        };

        let tokens = generate(&model, &[1, 5, 10], &params);
        assert!(!tokens.is_empty(), "should generate at least 1 token");
        assert!(tokens.len() <= 10);
        for t in &tokens {
            assert!((*t as usize) < config.vocab_size,
                "token {t} out of vocab range {}", config.vocab_size);
        }
    }

    #[test]
    fn test_generate_respects_stop_token() {
        let (_config, model) = load_test_model();
        let params = GenerateParams {
            max_tokens: 100,
            temperature: 0.0,
            stop_tokens: vec![0],
            ..Default::default()
        };

        let tokens = generate(&model, &[1], &params);
        assert!(!tokens.contains(&0));
    }

    /// Concurrent generation: two threads call generate() on the same model simultaneously.
    /// Each gets its own InferenceSession — no data races, no serialization.
    #[test]
    fn test_concurrent_generate() {
        use std::sync::Arc;

        let (_config, model) = load_test_model();
        let model = Arc::new(model);

        let params = Arc::new(GenerateParams {
            max_tokens: 5,
            temperature: 0.0,
            ..Default::default()
        });

        let handles: Vec<_> = (0..4)
            .map(|i| {
                let model = Arc::clone(&model);
                let params = Arc::clone(&params);
                std::thread::spawn(move || {
                    let tokens = generate(&model, &[1, 2, 3 + i as u32], &params);
                    assert!(!tokens.is_empty(), "thread {i}: expected tokens");
                })
            })
            .collect();

        for h in handles {
            h.join().expect("concurrent generate thread panicked");
        }
    }
}
