use anyhow::Result;
use meridian::llms::{
    open_ai::{messages::OpenAIMessage, OpenAIClient},
    LLMProvider, StructuredLLMProvider,
};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};

use crate::common::utils::{can_run_openai_test, setup_test_env, PromptFixture};

#[test]
fn test_openai_basic_completion() -> Result<()> {
    setup_test_env()?;

    let open_ai_user_messages = PromptFixture::open_ai_user_messages();
    let open_ai_client = OpenAIClient::new();

    for message in open_ai_user_messages {
        let completion = open_ai_client.get_completion(vec![message])?;
        assert!(!completion.is_empty());
    }

    Ok(())
}

#[test]
fn test_openai_structured_response() -> Result<()> {
    setup_test_env()?;

    if !can_run_openai_test() {
        return Ok(());
    }

    #[derive(Debug, Serialize, Deserialize, JsonSchema, PartialEq)]
    struct Step {
        explanation: String,
        output: String,
    }

    #[derive(Debug, Serialize, Deserialize, JsonSchema, PartialEq)]
    struct MathResponse {
        steps: Vec<Step>,
        final_answer: String,
    }

    let open_ai_client = OpenAIClient::new();
    let messages = vec![
        OpenAIMessage::System {
            content: "You are a helpful math tutor.".to_string(),
        },
        OpenAIMessage::User {
            content: "solve 8x + 31 = 2".to_string(),
        },
    ];

    let response: MathResponse =
        open_ai_client.get_structured_response::<MathResponse>(messages)?;

    // Validate the response
    assert!(!response.steps.is_empty(), "Should have at least one step");
    assert!(
        !response.final_answer.is_empty(),
        "Should have a final answer"
    );

    // Print the response for debugging
    println!("Structured Response: {:#?}", response);

    Ok(())
}