use crate::common::{setup_mock_llm, setup_mock_structured_llm};
use kernelx_core::prelude::*;

mod common;

#[tokio::test]
async fn test_llm_capabilities() {
    let provider = setup_mock_llm().unwrap();
    let model = provider.get_model::<dyn LLM>("mock-model").unwrap();

    let chat_result = model.chat(vec![ChatMessage::user("Hello")]).await;
    assert!(chat_result.is_ok());
    assert_eq!(chat_result.unwrap(), "Mock chat response");

    let complete_result = model.complete("Hello").await;
    assert!(complete_result.is_ok());
    assert_eq!(complete_result.unwrap(), "Mock completion response");
}

#[tokio::test]
async fn test_structured_llm_capabilities() {
    let provider = setup_mock_structured_llm().unwrap();
    let model = provider
        .get_model::<dyn StructuredLM>("mock-model")
        .unwrap();

    let schema = serde_json::json!({
        "type": "object",
        "properties": {
            "message": {"type": "string"},
            "confidence": {"type": "number"},
            "tags": {"type": "array", "items": {"type": "string"}}
        }
    });

    let structured_result = model.structured_complete("Test", &schema).await;
    assert!(structured_result.is_ok());

    let value = structured_result.unwrap();
    assert_eq!(value["message"], "Mock structured response");
    assert_eq!(value["confidence"], 0.95);
}