//! Test inference with long texts (8K, 16K, 32K tokens) on ModernBERT-base-32k
//!
//! This example tests inference with very long texts to verify 32K context support:
//! 1. Load ModernBERT-base-32k base model
//! 2. Load traditional classifier weights (intent, PII, jailbreak)
//! 3. Test inference on texts of various lengths (512, 8K, 16K, 32K tokens)
//! 4. Measure performance and verify accuracy
//!
//! Usage:
//!   cargo run --example test_long_text_inference_32k --release --no-default-features

use anyhow::{anyhow, Result};
use candle_core::{DType, Device, IndexOp, Tensor};
use candle_nn::{Linear, Module, VarBuilder};
use candle_semantic_router::model_architectures::traditional::modernbert::ModernBertVariant;
use candle_transformers::models::modernbert::{Config, ModernBert};
use hf_hub::{api::sync::Api, Repo, RepoType};
use std::path::Path;
use tokenizers::Tokenizer;

fn main() -> Result<()> {
    println!("Testing Long Text Inference with ModernBERT-base-32k");
    println!("{}", "=".repeat(70));

    let device = Device::Cpu; // Force CPU for testing (GPU would be faster)

    // Step 1: Download and load ModernBERT-base-32k
    println!("\nDownloading ModernBERT-base-32k base model...");
    let base_model_id = "llm-semantic-router/modernbert-base-32k";
    let repo = Repo::with_revision(
        base_model_id.to_string(),
        RepoType::Model,
        "main".to_string(),
    );
    let api = Api::new()?;
    let api = api.repo(repo);

    let base_config_path = api
        .get("config.json")
        .map_err(|e| anyhow!("Failed to download config.json: {}", e))?;
    let base_tokenizer_path = api
        .get("tokenizer.json")
        .map_err(|e| anyhow!("Failed to download tokenizer.json: {}", e))?;
    let base_weights_path = api
        .get("model.safetensors")
        .map_err(|e| anyhow!("Failed to download model.safetensors: {}", e))?;

    let base_model_dir = base_config_path.parent().unwrap();
    println!("   Base model directory: {:?}", base_model_dir);

    // Load base model config
    let config_str = std::fs::read_to_string(&base_config_path)?;
    let config: Config = serde_json::from_str(&config_str)?;
    println!(
        "   Config loaded: hidden_size={}, vocab_size={}",
        config.hidden_size, config.vocab_size
    );

    // Step 2: Load base model
    println!("\nLoading ModernBERT-base-32k base model...");
    let base_vb = unsafe {
        VarBuilder::from_mmaped_safetensors(&[base_weights_path.clone()], DType::F32, &device)
            .map_err(|e| anyhow!("Failed to load base model weights: {}", e))?
    };
    let base_model = ModernBert::load(base_vb, &config)
        .map_err(|e| anyhow!("Failed to load base ModernBert model: {}", e))?;
    println!("   Base model loaded successfully!");

    // Step 3: Load tokenizer
    println!("\nLoading tokenizer...");
    let mut tokenizer = Tokenizer::from_file(&base_tokenizer_path)
        .map_err(|e| anyhow!("Failed to load tokenizer: {}", e))?;

    // Configure padding
    if let Some(pad_token) = tokenizer.get_padding_mut() {
        pad_token.strategy = tokenizers::PaddingStrategy::BatchLongest;
        pad_token.pad_id = config.pad_token_id;
        pad_token.pad_token = ModernBertVariant::Extended32K.pad_token().to_string();
    }
    println!("   Tokenizer loaded and configured for 32K tokens");

    // Step 4: Test with different classifiers
    println!("\nTesting with different classifiers...");

    let classifiers = vec![
        (
            "Intent Classifier",
            "../models/lora_intent_classifier_bert-base-uncased_model",
            3,
        ),
        (
            "PII Detector",
            "../models/lora_pii_detector_bert-base-uncased_model",
            35,
        ),
        (
            "Jailbreak Classifier",
            "../models/lora_jailbreak_classifier_bert-base-uncased_model",
            2,
        ),
    ];

    for (classifier_name, classifier_path, num_classes) in classifiers {
        println!("\n{}", "=".repeat(60));
        println!("Testing: {}", classifier_name);
        println!("{}", "=".repeat(60));

        let model_weights_path = format!("{}/model.safetensors", classifier_path);
        if !Path::new(&model_weights_path).exists() {
            println!("   Model weights not found: {}", model_weights_path);
            continue;
        }

        // Load classifier
        let classifier_vb = unsafe {
            VarBuilder::from_mmaped_safetensors(&[model_weights_path.clone()], DType::F32, &device)
                .map_err(|e| anyhow!("Failed to load classifier weights: {}", e))?
        };

        let classifier_weight = classifier_vb
            .get((num_classes, config.hidden_size), "classifier.weight")
            .map_err(|e| anyhow!("Failed to load classifier weight: {}", e))?;
        let classifier_bias = classifier_vb.get(num_classes, "classifier.bias").ok();
        let classifier_head = Linear::new(classifier_weight, classifier_bias);
        println!("   Classifier loaded: {} classes", num_classes);

        // Step 5: Test with texts of various lengths
        println!("\nTesting inference with texts of various lengths...");

        // Create base text for repetition
        let base_text = "This is a test sentence for ModernBERT-base-32k integration. ";

        // Estimate tokens per character (roughly 1 token per 4 characters for English)
        let chars_per_token = 4.0;

        // Target token counts
        let target_tokens = vec![
            ("512 tokens", 512),
            ("1K tokens", 1024),
            ("8K tokens", 8192),
            ("16K tokens", 16384),
            ("32K tokens", 32768),
        ];

        for (name, target_token_count) in target_tokens {
            println!(
                "\n   Testing: {} (target: {} tokens)",
                name, target_token_count
            );

            // Calculate how many times to repeat base_text
            let chars_needed = (target_token_count as f64 * chars_per_token) as usize;
            let repetitions = chars_needed / base_text.len();
            let test_text = base_text.repeat(repetitions);

            let start = std::time::Instant::now();

            // Tokenize
            let encoding = tokenizer
                .encode(test_text.as_str(), true)
                .map_err(|e| anyhow!("Failed to encode text: {}", e))?;
            let input_ids: Vec<u32> = encoding.get_ids().to_vec();
            let attention_mask: Vec<u32> = encoding.get_attention_mask().to_vec();

            let actual_tokens = input_ids.len();
            println!("      Text length: {} characters", test_text.len());
            println!(
                "      Actual tokens: {} (target: {})",
                actual_tokens, target_token_count
            );

            // Skip if too long for CPU (32K tokens takes 10+ minutes)
            let is_cpu = matches!(device, Device::Cpu);
            if actual_tokens > 2000 && is_cpu {
                println!("      Skipping on CPU (would take too long)");
                println!("      Run on GPU for full 32K token testing");
                continue;
            }

            // Create tensors
            let input_ids_tensor = Tensor::new(&input_ids[..], &device)?.unsqueeze(0)?;
            let attention_mask_tensor = Tensor::new(&attention_mask[..], &device)?.unsqueeze(0)?;

            // Forward through base model
            let model_output = base_model
                .forward(&input_ids_tensor, &attention_mask_tensor)
                .map_err(|e| anyhow!("Base model forward failed: {}", e))?;

            // Pool: Use CLS token (first token)
            let pooled = model_output.i((.., 0, ..))?; // (batch_size, hidden_size)

            // Apply classifier
            let logits = classifier_head
                .forward(&pooled)
                .map_err(|e| anyhow!("Classifier forward failed: {}", e))?;

            // Apply softmax
            let probabilities = candle_nn::ops::softmax(&logits, 1)
                .map_err(|e| anyhow!("Softmax failed: {}", e))?;

            let probabilities_vec = probabilities.squeeze(0)?.to_vec1::<f32>()?;

            let (predicted_class, max_prob) = probabilities_vec
                .iter()
                .enumerate()
                .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
                .unwrap_or((0, &0.0));

            let elapsed = start.elapsed();
            println!("      Inference successful!");
            println!("         Predicted class: {}", predicted_class);
            println!("         Confidence: {:.4}", max_prob);
            println!(
                "         Latency: {:.2}ms ({:.2}s)",
                elapsed.as_secs_f64() * 1000.0,
                elapsed.as_secs_f64()
            );
            println!(
                "         Throughput: {:.2} tokens/sec",
                actual_tokens as f64 / elapsed.as_secs_f64()
            );
        }
    }

    println!("\nSummary:");
    println!("   ModernBERT-base-32k base model: Loaded");
    println!("   All classifiers: Loaded and compatible");
    println!("   Inference: Working on various text lengths");
    println!("   Note: Full 32K token testing requires GPU (CPU too slow)");
    println!("   All traditional classifiers are compatible with ModernBERT-base-32k!");

    println!("\nLong text inference test completed!");
    Ok(())
}
