baked-potato 0.4.0

use crate::error::MockError;
use mockito;
use potato_agent::agents::provider::{gemini::GenerateContentResponse, openai::OpenAIChatResponse};
use serde_json;

use pyo3::prelude::*;

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 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 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();

        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())
            .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(serde_json::to_string(&chat_structured_score_response).unwrap())
            .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(serde_json::to_string(&gemini_chat_response_with_score).unwrap())
            .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();

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

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

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

#[pymethods]
impl LLMTestServer {
    #[new]
    pub fn new() -> Self {
        LLMTestServer {
            openai_server: 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");
        }
        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(
                "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(),
            );
            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");

        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");
        Ok(())
    }

    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: PyObject,
        _exc_value: PyObject,
        _traceback: PyObject,
    ) -> Result<(), MockError> {
        self.stop_server()
    }
}

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