use sutra_afm::{FoundationModel, Session};
use serde::Deserialize;
use futures_util::StreamExt;

#[tokio::test]
async fn test_scenario_summarization() {
    let text = "The Open-Source Array Framework for Apple Silicon. MLX is an array framework for machine learning research on Apple silicon, brought to you by Apple machine learning research. MLX is designed by machine learning researchers for machine learning researchers. The framework is intended to be user-friendly, but still efficient to train and deploy models. The design of the framework itself is also conceptually simple. We intend to make it easy for researchers to extend and improve MLX with the goal of quickly exploring new ideas.";
    let prompt = format!("Summarize this in one sentence: {}", text);
    
    let result = FoundationModel::respond(&prompt).await.expect("Summarization failed");
    println!("[Scenario: Summarization] Output: {}", result);
    assert!(!result.is_empty());
}

#[tokio::test]
async fn test_scenario_proofreading() {
    let text = "This is a simple sentence with some errors. It has speling mistakes and grammer issues.";
    let prompt = format!("Proofread and fix this text, return only the corrected version: {}", text);
    
    let result = FoundationModel::respond(&prompt).await.expect("Proofreading failed");
    println!("[Scenario: Proofreading] Output: {}", result);
    assert!(result.contains("spelling") && result.contains("grammar"));
}

#[tokio::test]
async fn test_scenario_instruction_following() {
    let prompt = "Write a very short story (3 sentences) about a cat. \
                  Constraint 1: DO NOT use the word 'meow' or 'purr'. \
                  Constraint 2: Respond strictly in 3 numbered bullet points.";
    
    let result = FoundationModel::respond(prompt).await.expect("Instruction following failed");
    println!("[Scenario: Instruction Following] Output:\n{}", result);
    
    let result_lc = result.to_lowercase();
    assert!(!result_lc.contains("meow"));
    assert!(!result_lc.contains("purr"));
}

#[tokio::test]
async fn test_scenario_streaming() {
    let prompt = "Count from 1 to 3.";
    let mut stream = FoundationModel::stream(prompt);
    let mut full_response = String::new();
    
    while let Some(chunk) = stream.next().await {
        let chunk = chunk.expect("Stream chunk failed");
        full_response.push_str(&chunk);
    }
    assert!(full_response.contains("1") && full_response.contains("3"));
}

#[tokio::test]
async fn test_scenario_multi_turn() {
    let session = Session::new("test_async_session", "You are a helpful assistant.");
    let _ = session.respond("My favorite color is blue.").await.unwrap();
    let res = session.respond("What is my favorite color?").await.unwrap();
    println!("[Scenario: Multi-turn] Response: {}", res);
    assert!(res.to_lowercase().contains("blue"));
}

#[derive(Deserialize, Debug)]
struct SentimentResponse {
    sentiment: String,
    confidence: f32,
}

#[tokio::test]
async fn test_scenario_sentiment_analysis() {
    let text = "I absolutely love the new M4 Mac Mini!";
    let prompt = format!("Analyze sentiment of: '{}'. JSON fields: 'sentiment', 'confidence'.", text);
    
    let result = FoundationModel::respond_json::<SentimentResponse>(&prompt).await.expect("Sentiment analysis failed");
    println!("[Scenario: Sentiment Analysis] Output: {}, Confidence: {}", result.sentiment, result.confidence);
    assert_eq!(result.sentiment.to_lowercase(), "positive");
}
