gemini_bridge 0.1.15

use anyhow::Error;
use async_trait::async_trait;
use serde_json::json;

use super::{request::RequestBody, response::GenerateContentResponse};

fn build_generate_content_uri(api_key: String, model: String) -> String {
    let result = format!(
        "https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent?key={}",
        model, api_key
    );
    return result;
}

// fn build_generate_stream_uri(api_key: String, model: String) -> String {
//     let result = format!(
//         "https://generativelanguage.googleapis.com/v1beta/models/{}:streamGenerateContent?alt=sse&key={}",
//         model, api_key
//     );
//     return result;
// }

#[async_trait]
pub trait TextGenerator: Send + Sync {
    async fn generate_content(&self, req: RequestBody) -> Result<GenerateContentResponse, Error>;
    // async fn generate_content_stream(
    //     &self,
    //     req: RequestBody,
    // ) -> Result<GenerateContentResponse, Error>;
}

pub(crate) struct TextGeneratorImpl {
    pub api_key: String,
    pub model: String,
}

#[async_trait]
impl TextGenerator for TextGeneratorImpl {
    async fn generate_content(&self, req: RequestBody) -> Result<GenerateContentResponse, Error> {
        let uri = &build_generate_content_uri(self.api_key.to_owned(), self.model.to_owned());

        let request_body = json!(req);

        let res = reqwest::Client::new()
            .post(uri)
            .json(&request_body)
            .send()
            .await;

        let response_body = match res {
            Ok(response) => {
                if response.status().is_success() {
                    response.text().await?
                } else {
                    return Err(Error::msg(format!(
                        "Request failed with status: {}",
                        response.status()
                    )));
                }
            }
            Err(_) => return Err(Error::msg("Failed to send request")),
        };

        // Attempt to deserialize the response body
        match serde_json::from_str::<GenerateContentResponse>(&response_body) {
            Ok(parsed_response) => Ok(parsed_response),
            Err(e) => {
                println!("Failed to parse response body: {}", e);
                println!("Response body: {}", response_body);
                Err(Error::msg("Failed to parse response body"))
            }
        }
    }
}