baked-potato 0.26.0

use crate::error::MockError;
use mockito;
use potato_type::anthropic::AnthropicMessageResponse;
use potato_type::google::v1::generate::{DataNum, GenerateContentResponse};
use potato_type::google::GeminiEmbeddingResponse;
use potato_type::openai::v1::embedding::OpenAIEmbeddingResponse;
use potato_type::openai::v1::OpenAIChatResponse;
use potato_type::openai::{ChatMessage, ContentPart, TextContentPart};
use potato_type::prompt::{MessageNum, Prompt, Role, Score};
use potato_type::StructuredOutput;
use pyo3::prelude::*;
use rand::Rng;
use serde_json;

pub const OPENAI_EMBEDDING_RESPONSE: &str = include_str!("assets/openai/embedding_response.json");

pub const GEMINI_EMBEDDING_RESPONSE: &str = include_str!("assets/gemini/embedding_response.json");

pub const OPENAI_CHAT_COMPLETION_RESPONSE: &str =
    include_str!("assets/openai/openai_chat_completion_response.json");

pub const OPENAI_CHAT_STRUCTURED_RESPONSE: &str =
    include_str!("assets/openai/chat_completion_structured_response.json");

pub const OPENAI_CHAT_STRUCTURED_SCORE_RESPONSE: &str =
    include_str!("assets/openai/chat_completion_structured_score_response.json");

pub const OPENAI_CHAT_STRUCTURED_RESPONSE_PARAMS: &str =
    include_str!("assets/openai/chat_completion_structured_response_params.json");

pub const OPENAI_CHAT_STRUCTURED_TASK_OUTPUT: &str =
    include_str!("assets/openai/chat_completion_structured_task_output.json");

pub const GEMINI_CHAT_COMPLETION_RESPONSE: &str =
    include_str!("assets/gemini/chat_completion.json");

pub const GEMINI_CHAT_COMPLETION_RESPONSE_WITH_SCORE: &str =
    include_str!("assets/gemini/chat_completion_with_score.json");

pub const ANTHROPIC_MESSAGE_RESPONSE: &str =
    include_str!("assets/anthropic/message_completion.json");

pub const ANTHROPIC_MESSAGE_STRUCTURED_RESPONSE: &str =
    include_str!("assets/anthropic/message_structured_completion.json");

pub const ANTHROPIC_MESSAGE_STRUCTURED_TASK_OUTPUT: &str =
    include_str!("assets/anthropic/message_structured_completion_tasks.json");

pub const OPENAI_CHAT_TOOL_CALL_RESPONSE: &str =
    include_str!("assets/openai/chat_completion_tool_call.json");

pub const GEMINI_CHAT_TOOL_CALL_RESPONSE: &str =
    include_str!("assets/gemini/chat_completion_tool_call.json");

pub const ANTHROPIC_MESSAGE_TOOL_USE_RESPONSE: &str =
    include_str!("assets/anthropic/message_tool_use.json");

fn randomize_openai_embedding_response(
    response: OpenAIEmbeddingResponse,
) -> OpenAIEmbeddingResponse {
    // create random Vec<f32> of length 512
    let mut cloned_response = response.clone();
    let mut rng = rand::rng();
    let embedding: Vec<f32> = (0..512).map(|_| rng.random_range(-1.0..1.0)).collect();
    cloned_response.data[0].embedding = embedding;
    cloned_response
}

fn randomize_gemini_embedding_response(
    response: GeminiEmbeddingResponse,
) -> GeminiEmbeddingResponse {
    let mut cloned_response = response.clone();
    let mut rng = rand::rng();
    let embedding: Vec<f32> = (0..512).map(|_| rng.random_range(-1.0..1.0)).collect();
    cloned_response.embedding.values = embedding;
    cloned_response
}

fn randomize_structured_openai_score_response(response: &OpenAIChatResponse) -> OpenAIChatResponse {
    let mut cloned_response = response.clone();
    let mut rng = rand::rng();

    // Generate random score between 1 and 5
    let score = rng.random_range(1..=5);

    // Generate random reason from a set of predefined reasons
    let reasons = [
        "The code is excellent and follows best practices.",
        "The implementation is solid with minor improvements possible.",
        "The code works but could use some optimization.",
        "The solution is functional but needs refactoring.",
        "The code has significant issues that need addressing.",
    ];
    let reason = reasons[rng.random_range(0..reasons.len())];

    cloned_response.choices[0].message.content = Some(format!(
        "{{ \"score\": {}, \"reason\": \"{}\" }}",
        score, reason
    ));

    cloned_response
}

fn randomize_gemini_score_response(response: GenerateContentResponse) -> GenerateContentResponse {
    let mut cloned_response = response.clone();
    let mut rng = rand::rng();

    // Generate random score between 1 and 5 (typical for Gemini scoring)
    let score = rng.random_range(1..=5);

    // Generate random reason from a set of predefined reasons
    let reasons = [
        "The model performed exceptionally well on the evaluation.",
        "Good performance with room for minor improvements.",
        "Satisfactory results with some areas for optimization.",
        "Adequate performance but needs significant improvements.",
        "Performance below expectations, requires major adjustments.",
    ];
    let reason = reasons[rng.random_range(0..reasons.len())];

    // Update the first candidate's content
    if let Some(candidate) = cloned_response.candidates.get_mut(0) {
        if let Some(part) = candidate.content.parts.get_mut(0) {
            part.data = DataNum::Text(format!(
                "{{\"score\": {}, \"reason\": \"{}\"}}",
                score, reason
            ));
        }
    }

    cloned_response
}

pub struct LLMApiMock {
    pub url: String,
    pub server: mockito::ServerGuard,
}

impl LLMApiMock {
    pub fn new() -> Self {
        let mut server = mockito::Server::new();
        // load the OpenAI chat completion response
        let openai_embedding_response: OpenAIEmbeddingResponse =
            serde_json::from_str(OPENAI_EMBEDDING_RESPONSE).unwrap();
        let chat_msg_response: OpenAIChatResponse =
            serde_json::from_str(OPENAI_CHAT_COMPLETION_RESPONSE).unwrap();
        let chat_structured_response: OpenAIChatResponse =
            serde_json::from_str(OPENAI_CHAT_STRUCTURED_RESPONSE).unwrap();
        let chat_structured_score_response: OpenAIChatResponse =
            serde_json::from_str(OPENAI_CHAT_STRUCTURED_SCORE_RESPONSE).unwrap();
        let chat_structured_response_params: OpenAIChatResponse =
            serde_json::from_str(OPENAI_CHAT_STRUCTURED_RESPONSE_PARAMS).unwrap();
        let chat_structured_task_output: OpenAIChatResponse =
            serde_json::from_str(OPENAI_CHAT_STRUCTURED_TASK_OUTPUT).unwrap();

        // load the Gemini chat completion response
        let gemini_chat_response: GenerateContentResponse =
            serde_json::from_str(GEMINI_CHAT_COMPLETION_RESPONSE).unwrap();
        let gemini_chat_response_with_score: GenerateContentResponse =
            serde_json::from_str(GEMINI_CHAT_COMPLETION_RESPONSE_WITH_SCORE).unwrap();
        let gemini_embedding_response: GeminiEmbeddingResponse =
            serde_json::from_str(GEMINI_EMBEDDING_RESPONSE).unwrap();

        // anthropic message response
        let anthropic_message_response: AnthropicMessageResponse =
            serde_json::from_str(ANTHROPIC_MESSAGE_RESPONSE).unwrap();

        let anthropic_message_structured_response: AnthropicMessageResponse =
            serde_json::from_str(ANTHROPIC_MESSAGE_STRUCTURED_RESPONSE).unwrap();

        let anthropic_message_structured_task_output: AnthropicMessageResponse =
            serde_json::from_str(ANTHROPIC_MESSAGE_STRUCTURED_TASK_OUTPUT).unwrap();

        server
            .mock("POST", "/chat/completions")
            .match_body(mockito::Matcher::PartialJson(serde_json::json!({
                "response_format": {
                    "type": "json_schema",
                    "json_schema": {
                        "name": "Parameters",
                         "schema": {
                              "$schema": "https://json-schema.org/draft/2020-12/schema",
                              "properties": {
                                  "variable1": {
                                  "format": "int32",
                                  "type": "integer"
                                  },
                                  "variable2": {
                                  "format": "int32",
                                  "type": "integer"
                                  }
                              },
                              "required": [
                                  "variable1",
                                  "variable2"
                              ],
                              "title": "Parameters",
                              "type": "object"
                              },
                        "strict": true
                    }

                }
            })))
            .expect(usize::MAX)
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body(serde_json::to_string(&chat_structured_response_params).unwrap())
            .create();

        server
            .mock("POST", "/chat/completions")
            .match_body(mockito::Matcher::PartialJson(serde_json::json!({
               "response_format": {
                    "type": "json_schema",
                    "json_schema": {
                        "name": "TaskOutput",
                    }
                }
            })))
            .expect(usize::MAX)
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body(serde_json::to_string(&chat_structured_task_output).unwrap())
            .create();

        server
            .mock("POST", "/chat/completions")
            .match_body(mockito::Matcher::PartialJson(serde_json::json!({
                "response_format": {
                    "type": "json_schema",
                    "json_schema": {
                        "name": "Score",
                    }
                }
            })))
            .expect(usize::MAX)
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body(serde_json::to_string(&chat_structured_score_response).unwrap())
            .with_body_from_request({
                let chat_structured_score_response = chat_structured_score_response.clone();
                move |_request| {
                    let randomized_response = randomize_structured_openai_score_response(
                        &chat_structured_score_response.clone(),
                    );
                    serde_json::to_string(&randomized_response).unwrap().into()
                }
            })
            .create();

        server
            .mock("POST", "/chat/completions")
            .match_body(mockito::Matcher::Regex(
                r#".*"name"\s*:\s*"Score".*"#.to_string(),
            ))
            .expect(usize::MAX)
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body_from_request({
                let chat_structured_score_response = chat_structured_score_response.clone();
                move |_request| {
                    let randomized_response = randomize_structured_openai_score_response(
                        &chat_structured_score_response.clone(),
                    );
                    serde_json::to_string(&randomized_response).unwrap().into()
                }
            })
            .create();

        server
            .mock("POST", "/chat/completions")
            .match_body(mockito::Matcher::PartialJson(serde_json::json!({
                "response_format": {
                    "type": "json_schema"
                }
            })))
            .expect(usize::MAX)
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body(serde_json::to_string(&chat_structured_response).unwrap())
            .create();

        // mock the Gemini chat completion response
        server
            .mock(
                "POST",
                mockito::Matcher::Regex(r".*/.*:generateContent$".to_string()),
            )
            .match_header("x-goog-api-key", mockito::Matcher::Any)
            .match_header("content-type", "application/json")
            .match_body(mockito::Matcher::PartialJson(serde_json::json!({
                "contents": [
                    {
                        "parts": [
                            {
                                "text":  "You are a helpful assistant"
                            }
                        ]
                    }
                ]
            })))
            .expect(usize::MAX) // More specific expectation than usize::MAX
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body(serde_json::to_string(&gemini_chat_response).unwrap())
            .create();

        // mock structured response
        server
            .mock(
                "POST",
                mockito::Matcher::Regex(r".*/.*:generateContent$".to_string()),
            )
            .match_header("x-goog-api-key", mockito::Matcher::Any)
            .match_header("content-type", "application/json")
            .match_body(mockito::Matcher::PartialJson(serde_json::json!({
                "generation_config": {
                    "responseMimeType": "application/json"
                }
            })))
            .expect(usize::MAX)
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body_from_request(move |_request| {
                let randomized_response =
                    randomize_gemini_score_response(gemini_chat_response_with_score.clone());
                serde_json::to_string(&randomized_response).unwrap().into()
            })
            .create();

        // Openai chat completion mock
        server
            .mock("POST", "/chat/completions")
            .expect(usize::MAX)
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body(serde_json::to_string(&chat_msg_response).unwrap())
            .create();

        server
            .mock("POST", "/embeddings")
            .expect(usize::MAX)
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body_from_request(move |_request| {
                let randomized_response =
                    randomize_openai_embedding_response(openai_embedding_response.clone());
                serde_json::to_string(&randomized_response).unwrap().into()
            })
            .create();

        server
            .mock(
                "POST",
                mockito::Matcher::Regex(r".*/.*:embedContent$".to_string()),
            )
            .expect(usize::MAX)
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body_from_request(move |_request| {
                let randomized_response =
                    randomize_gemini_embedding_response(gemini_embedding_response.clone());
                serde_json::to_string(&randomized_response).unwrap().into()
            })
            .create();

        // mock the anthropic message response

        server
            .mock("POST", "/messages")
            .match_header("content-type", "application/json")
            .match_body(mockito::Matcher::PartialJson(serde_json::json!({
                "messages": [
                    {
                        "content": [
                            {
                                "text":  "Give me a score!",
                                "type": "text"
                            }
                        ]
                    }
                ]
            })))
            .expect(usize::MAX) // More specific expectation than usize::MAX
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body(serde_json::to_string(&anthropic_message_structured_response).unwrap())
            .create();

        server
            .mock("POST", "/messages")
            .match_header("content-type", "application/json")
            .match_body(mockito::Matcher::Regex(
                r#".*"text"\s*:\s*"Give me a task list!".*"#.to_string(),
            ))
            .expect(usize::MAX) // More specific expectation than usize::MAX
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body(serde_json::to_string(&anthropic_message_structured_task_output).unwrap())
            .create();

        server
            .mock("POST", "/messages")
            .expect(usize::MAX)
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body(serde_json::to_string(&anthropic_message_response).unwrap())
            .create();

        Self {
            url: server.url(),
            server,
        }
    }
}

impl LLMApiMock {
    /// Register tool-call response mocks for all three providers.
    ///
    /// Mockito's last-registered-wins matching ensures the tool-call mock fires first
    /// (one shot via `.expect(1)`), then subsequent requests fall through to the
    /// existing text-response mock.
    pub fn enable_tool_call_flow(&mut self) {
        let openai_tool_response: OpenAIChatResponse =
            serde_json::from_str(OPENAI_CHAT_TOOL_CALL_RESPONSE).unwrap();
        let gemini_tool_response: GenerateContentResponse =
            serde_json::from_str(GEMINI_CHAT_TOOL_CALL_RESPONSE).unwrap();
        let gemini_text_response: GenerateContentResponse =
            serde_json::from_str(GEMINI_CHAT_COMPLETION_RESPONSE).unwrap();
        let anthropic_tool_response: AnthropicMessageResponse =
            serde_json::from_str(ANTHROPIC_MESSAGE_TOOL_USE_RESPONSE).unwrap();

        // OpenAI tool call mock — fires once then falls through to the text response mock
        self.server
            .mock("POST", "/chat/completions")
            .expect(1)
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body(serde_json::to_string(&openai_tool_response).unwrap())
            .create();

        // Gemini catch-all text response — lower priority (registered before the tool call mock),
        // fires for any Gemini request once the one-shot tool call mock is consumed.
        self.server
            .mock(
                "POST",
                mockito::Matcher::Regex(r".*/.*:generateContent$".to_string()),
            )
            .match_header("x-goog-api-key", mockito::Matcher::Any)
            .match_header("content-type", "application/json")
            .expect(usize::MAX)
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body(serde_json::to_string(&gemini_text_response).unwrap())
            .create();

        // Gemini tool call mock — fires once (registered last = highest priority)
        self.server
            .mock(
                "POST",
                mockito::Matcher::Regex(r".*/.*:generateContent$".to_string()),
            )
            .match_header("x-goog-api-key", mockito::Matcher::Any)
            .match_header("content-type", "application/json")
            .expect(1)
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body(serde_json::to_string(&gemini_tool_response).unwrap())
            .create();

        // Anthropic tool use mock — fires once
        self.server
            .mock("POST", "/messages")
            .match_header("content-type", "application/json")
            .expect(1)
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body(serde_json::to_string(&anthropic_tool_response).unwrap())
            .create();
    }
}

impl Default for LLMApiMock {
    fn default() -> Self {
        Self::new()
    }
}

#[pyclass(skip_from_py_object)]
#[allow(dead_code)]
pub struct LLMTestServer {
    openai_server: Option<LLMApiMock>,

    #[pyo3(get)]
    pub url: Option<String>,
}

#[pymethods]
impl LLMTestServer {
    #[new]
    pub fn new() -> Self {
        LLMTestServer {
            openai_server: None,
            url: None,
        }
    }

    pub fn start_mock_server(&mut self) -> Result<(), MockError> {
        let llm_server = LLMApiMock::new();
        println!("Mock LLM Server started at {}", llm_server.url);
        self.openai_server = Some(llm_server);
        Ok(())
    }

    pub fn stop_mock_server(&mut self) {
        if let Some(server) = self.openai_server.take() {
            drop(server);
            std::env::remove_var("OPENAI_API_URL");
            std::env::remove_var("OPENAI_API_KEY");
            std::env::remove_var("GEMINI_API_KEY");
            std::env::remove_var("GEMINI_API_URL");
            std::env::remove_var("GOOGLE_API_KEY");
            std::env::remove_var("GOOGLE_API_URL");
            std::env::remove_var("ANTHROPIC_API_KEY");
            std::env::remove_var("ANTHROPIC_API_URL");
        }
        println!("Mock LLM Server stopped");
    }

    pub fn set_env_vars_for_client(&self) -> Result<(), MockError> {
        {
            std::env::set_var("APP_ENV", "dev_client");
            std::env::set_var("OPENAI_API_KEY", "test_key");
            std::env::set_var("GEMINI_API_KEY", "gemini");
            std::env::set_var("GOOGLE_API_KEY", "google");
            std::env::set_var("ANTHROPIC_API_KEY", "anthropic_key");
            std::env::set_var(
                "OPENAI_API_URL",
                self.openai_server.as_ref().unwrap().url.clone(),
            );
            std::env::set_var(
                "GEMINI_API_URL",
                self.openai_server.as_ref().unwrap().url.clone(),
            );
            std::env::set_var(
                "GOOGLE_API_URL",
                self.openai_server.as_ref().unwrap().url.clone(),
            );
            std::env::set_var(
                "ANTHROPIC_API_URL",
                self.openai_server.as_ref().unwrap().url.clone(),
            );

            Ok(())
        }
    }

    pub fn start_server(&mut self) -> Result<(), MockError> {
        self.cleanup()?;

        println!("Starting Mock GenAI Server...");
        self.start_mock_server()?;
        self.set_env_vars_for_client()?;

        // set server env vars
        std::env::set_var("APP_ENV", "dev_server");

        self.url = Some(self.openai_server.as_ref().unwrap().url.clone());

        Ok(())
    }

    pub fn stop_server(&mut self) -> Result<(), MockError> {
        self.cleanup()?;

        Ok(())
    }

    pub fn remove_env_vars_for_client(&self) -> Result<(), MockError> {
        std::env::remove_var("OPENAI_API_URI");
        std::env::remove_var("OPENAI_API_KEY");
        std::env::remove_var("GEMINI_API_KEY");
        std::env::remove_var("GEMINI_API_URL");
        std::env::remove_var("GOOGLE_API_KEY");
        std::env::remove_var("GOOGLE_API_URL");
        std::env::remove_var("ANTHROPIC_API_KEY");
        std::env::remove_var("ANTHROPIC_API_URL");
        Ok(())
    }

    /// Enable tool call flow on the underlying mock server.
    /// Must be called after `start_server()`.
    pub fn enable_tool_call_flow(&mut self) {
        if let Some(ref mut mock) = self.openai_server {
            mock.enable_tool_call_flow();
        }
    }

    fn cleanup(&self) -> Result<(), MockError> {
        // unset env vars
        self.remove_env_vars_for_client()?;

        Ok(())
    }

    fn __enter__(mut self_: PyRefMut<Self>) -> Result<PyRefMut<Self>, MockError> {
        self_.start_server()?;

        Ok(self_)
    }

    fn __exit__(
        &mut self,
        _exc_type: Py<PyAny>,
        _exc_value: Py<PyAny>,
        _traceback: Py<PyAny>,
    ) -> Result<(), MockError> {
        self.stop_server()
    }
}

impl Default for LLMTestServer {
    fn default() -> Self {
        Self::new()
    }
}

#[allow(clippy::uninlined_format_args)]
pub fn create_score_prompt(params: Option<Vec<String>>) -> Prompt {
    let mut user_prompt = "What is the score?".to_string();

    // If parameters are provided, append them to the user prompt in format ${param}
    if let Some(params) = params {
        for param in params {
            user_prompt.push_str(&format!(" ${{{}}}", param));
        }
    }

    let system_content = "You are a helpful assistant.".to_string();

    let system_msg = ChatMessage {
        role: Role::Developer.to_string(),
        content: vec![ContentPart::Text(TextContentPart::new(system_content))],
        name: None,
    };

    let user_msg = ChatMessage {
        role: Role::User.to_string(),
        content: vec![ContentPart::Text(TextContentPart::new(user_prompt))],
        name: None,
    };
    Prompt::new_rs(
        vec![MessageNum::OpenAIMessageV1(user_msg)],
        "gpt-4o",
        potato_type::Provider::OpenAI,
        vec![MessageNum::OpenAIMessageV1(system_msg)],
        None,
        Some(Score::get_structured_output_schema()),
        potato_type::prompt::ResponseType::Score,
    )
    .unwrap()
}

pub fn create_parameterized_prompt() -> Prompt {
    let user_content = "What is ${variable1} + ${variable2}?".to_string();
    let system_content = "You are a helpful assistant.".to_string();

    let system_msg = ChatMessage {
        role: Role::Developer.to_string(),
        content: vec![ContentPart::Text(TextContentPart::new(system_content))],
        name: None,
    };

    let user_msg = ChatMessage {
        role: Role::User.to_string(),
        content: vec![ContentPart::Text(TextContentPart::new(user_content))],
        name: None,
    };
    Prompt::new_rs(
        vec![MessageNum::OpenAIMessageV1(user_msg)],
        "gpt-4o",
        potato_type::Provider::OpenAI,
        vec![MessageNum::OpenAIMessageV1(system_msg)],
        None,
        None,
        potato_type::prompt::ResponseType::Null,
    )
    .unwrap()
}