use crate::common::setup_mock_structured_llm;
use kernelx_core::prelude::*;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};

mod common;

#[derive(Debug, Serialize, Deserialize, JsonSchema)]
struct TestResponse {
    message: String,
    confidence: f32,
    tags: Vec<String>,
}

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

    let schema = schemars::schema_for!(TestResponse);
    let result = model
        .structured_complete(
            "Generate a test response",
            &serde_json::to_value(schema.schema).unwrap(),
        )
        .await;

    assert!(result.is_ok());

    let value = result.unwrap();
    assert_eq!(value["message"], "Mock structured response");
    assert_eq!(value["confidence"], 0.95);
    assert!(value["tags"].as_array().unwrap().len() == 3);
}