//! OpenAI integration for [Flows.network](https://flows.network)
//!
//! # Quick Start
//!
//! To get started, let's write a very tiny flow function.
//!
//! ```rust
//! use openai_flows::{chat_completion, ChatModel, ChatOptions, FlowsAccount};
//! use lambda_flows::{request_received, send_response};
//! use serde_json::Value;
//! use std::collections::HashMap;
//!
//! #[no_mangle]
//! #[tokio::main(flavor = "current_thread")]
//! pub async fn run() {
//!     request_received(handler).await;
//! }
//!
//! async fn handler(_qry: HashMap<String, Value>, body: Vec<u8>) {
//!     let co = ChatOptions {
//!         model: ChatModel::GPT35Turbo,
//!         restart: false,
//!         system_prompt: None,
//!         retry_times: 2,
//!     };
//!     let r = match chat_completion(
//!         FlowsAccount::Default,
//!         "any_conversation_id",
//!         String::from_utf8_lossy(&body).into_owned().as_str(),
//!         &co,
//!     )
//!     .await
//!     {
//!         Ok(c) => c.choice,
//!         Err(e) => e,
//!     };
//!     
//!     send_response(
//!         200,
//!         vec![(
//!             String::from("content-type"),
//!             String::from("text/plain; charset=UTF-8"),
//!         )],
//!         r.as_bytes().to_vec(),
//!     );
//! }
//! ```
//!
//! When the Lambda request is received, chat
//! using [chat_completion] then send the response.

use std::fmt;
use std::thread::sleep;
use std::time::Duration;

use http_req::{
    request::{Method, Request},
    uri::Uri,
};
use lazy_static::lazy_static;
use serde::{Deserialize, Serialize};
use urlencoding::encode;

lazy_static! {
    static ref OPENAI_API_PREFIX: String = String::from(
        std::option_env!("OPENAI_API_PREFIX").unwrap_or("https://openai.flows.network/api")
    );
}

extern "C" {
    fn get_flows_user(p: *mut u8) -> i32;
    fn get_flow_id(p: *mut u8) -> i32;
}

const MAX_RETRY_TIMES: u8 = 10;
const RETRY_INTERVAL: u64 = 10; // Wait 10 seconds before retry

/// The account name you provide to
/// [Flows.network](https://flows.network) platform,
/// which is tied to your OpenAI API key.
///
/// If set as `Default`, the 'Default' named account will be used.
/// If there is no 'Default' named account,
/// a non-fixed one will be selected from all your connected accounts.
///
pub enum FlowsAccount {
    Default,
    Provided(String),
}

/// Request struct for the completion.
///
/// The default model is "text-davinci-003".
/// Use retry_times to set the number of retries when requesting
/// OpenAI's api encounters a problem. Default is 2 and max number is 10.
/// For more detail about parameters, please refer to
/// [OpenAI docs](https://platform.openai.com/docs/api-reference/completions/create)
///
#[derive(Debug, Serialize)]
pub struct CompletionRequest {
    pub model: String,
    pub prompt: String,
    pub suffix: Option<String>,
    pub n: u8,
    pub best_of: u8,
    pub max_tokens: u16,
    pub temperature: f32,
    pub top_p: f32,
    pub logprobs: Option<u8>,
    pub presence_penalty: f32,
    pub frequency_penalty: f32,
    #[serde(skip_serializing)]
    pub retry_times: u8,
}

impl Default for CompletionRequest {
    fn default() -> CompletionRequest {
        CompletionRequest {
            model: String::from("text-davinci-003"),
            prompt: String::from("<|endoftext|>"),
            suffix: None,
            n: 1,
            best_of: 1,
            max_tokens: 16,
            temperature: 1.0,
            top_p: 1.0,
            logprobs: None,
            presence_penalty: 0.0,
            frequency_penalty: 0.0,
            retry_times: 2,
        }
    }
}

/// Create completion for the provided prompt and parameters.
///
/// `account` is an [FlowsAccount] used for picking your tied OpenAI API key.
///
/// `params` is a [CompletionRequest] object.
///
/// If you have not connected your OpenAI account with [Flows.network platform](https://flows.network),
/// you will receive an error in the flow's building log or running log.
///
pub async fn create_completion(
    account: FlowsAccount,
    params: CompletionRequest,
) -> Result<Vec<String>, String> {
    let retry_times = match params.retry_times {
        r if r <= 0 => 1,
        r if r > MAX_RETRY_TIMES => MAX_RETRY_TIMES,
        r => r,
    };
    let account = match account {
        FlowsAccount::Default => String::new(),
        FlowsAccount::Provided(s) => s,
    };
    create_completion_inner(account, params, retry_times)
}

fn create_completion_inner(
    account: String,
    params: CompletionRequest,
    retry_times: u8,
) -> Result<Vec<String>, String> {
    unsafe {
        let mut flows_user = Vec::<u8>::with_capacity(100);
        let c = get_flows_user(flows_user.as_mut_ptr());
        flows_user.set_len(c as usize);
        let flows_user = String::from_utf8(flows_user).unwrap();

        let mut writer = Vec::new();
        let uri = format!(
            "{}/{}/create_completion?account={}",
            OPENAI_API_PREFIX.as_str(),
            flows_user,
            encode(account.as_str()),
        );
        let uri = Uri::try_from(uri.as_str()).unwrap();
        let body = serde_json::to_vec(&params).unwrap_or_default();
        match Request::new(&uri)
            .method(Method::POST)
            .header("Content-Type", "application/json")
            .header("Content-Length", &body.len())
            .body(&body)
            .send(&mut writer)
        {
            Ok(res) => {
                match res.status_code().is_success() {
                    true => serde_json::from_slice::<Vec<String>>(&writer)
                        .or(Err(String::from("Unexpected error"))),
                    false => {
                        match res.status_code().into() {
                            409 | 429 | 503 => {
                                // 409 TryAgain 429 RateLimitError
                                // 503 ServiceUnavableila
                                if retry_times > 1 {
                                    sleep(Duration::from_secs(RETRY_INTERVAL));
                                    return create_completion_inner(
                                        account,
                                        params,
                                        retry_times - 1,
                                    );
                                }
                            }
                            _ => {}
                        }
                        Err(String::from_utf8_lossy(&writer).into_owned())
                    }
                }
            }
            Err(e) => Err(e.to_string()),
        }
    }
}

/// Response struct for the chat completion.
///
/// `restarted` is the flag to show whether a new conversation is created.
///
/// `choice` is the response from ChatGPT.
///
#[derive(Debug, Deserialize)]
pub struct ChatResponse {
    pub restarted: bool,
    pub choice: String,
}

impl Default for ChatResponse {
    fn default() -> ChatResponse {
        ChatResponse {
            restarted: true,
            choice: String::new(),
        }
    }
}

/// Models for Chat
#[derive(Debug, Clone, Copy)]
pub enum ChatModel {
    GPT4_32K,
    GPT4,
    GPT35Turbo,
}

impl fmt::Display for ChatModel {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            ChatModel::GPT4_32K => write!(f, "gpt-4-32k"),
            ChatModel::GPT4 => write!(f, "gpt-4"),
            ChatModel::GPT35Turbo => write!(f, "gpt-3.5-turbo"),
        }
    }
}

impl Default for ChatModel {
    fn default() -> ChatModel {
        ChatModel::GPT35Turbo
    }
}

/// struct for setting the chat options.
///
/// When `restart` is true, a new conversation will be created.
///
/// Use retry_times to set the number of retries when requesting
/// OpenAI's api encounters a problem. The max number is 10.
///
/// `system_prompt` will be treated as the prompt of the system role.
///
#[derive(Debug, Default)]
pub struct ChatOptions<'a> {
    pub model: ChatModel,
    pub restart: bool,
    pub system_prompt: Option<&'a str>,
    pub retry_times: u8,
}

/// Create chat completion with the provided sentence.
/// It use OpenAI's [GPT-3.5](https://platform.openai.com/docs/models/gpt-3-5) model to make a conversation.
/// It will keep the conversation history for 10 minutes for you.
/// That means a new conversation will be created if you haven't call this method for 10 minutes.
///
/// `account` is an [FlowsAccount] used for picking your tied OpenAI API key.
///
/// `conversation_id` is the identity of the conversation. The history will be fetched and attached
/// to the `sentence` as a whole prompt for ChatGPT.
///
/// `sentence` is a String reprensent the sentence of the conversation.
///
/// If you have not connected your OpenAI account with [Flows.network platform](https://flows.network),
/// you will receive an error in the flow's building log or running log.
///
pub async fn chat_completion(
    account: FlowsAccount,
    conversation_id: &str,
    sentence: &str,
    options: &ChatOptions<'_>,
) -> Result<ChatResponse, String> {
    let retry_times = match options.retry_times {
        r if r <= 0 => 1,
        r if r > MAX_RETRY_TIMES => MAX_RETRY_TIMES,
        r => r,
    };
    let account = match account {
        FlowsAccount::Default => String::new(),
        FlowsAccount::Provided(s) => s,
    };
    chat_completion_inner(account, conversation_id, sentence, options, retry_times)
}

fn chat_completion_inner(
    account: String,
    conversation_id: &str,
    sentence: &str,
    options: &ChatOptions,
    retry_times: u8,
) -> Result<ChatResponse, String> {
    unsafe {
        let mut flows_user = Vec::<u8>::with_capacity(100);
        let c = get_flows_user(flows_user.as_mut_ptr());
        flows_user.set_len(c as usize);
        let flows_user = String::from_utf8(flows_user).unwrap();

        let mut flow_id = Vec::<u8>::with_capacity(100);
        let c = get_flow_id(flow_id.as_mut_ptr());
        if c == 0 {
            panic!("Failed to get flow id");
        }
        flow_id.set_len(c as usize);
        let flow_id = String::from_utf8(flow_id).unwrap();

        let mut writer = Vec::new();
        let uri = format!(
            "{}/{}/{}/chat_completion?account={}&conversation={}&model={}&restart={}",
            OPENAI_API_PREFIX.as_str(),
            flows_user,
            flow_id,
            encode(account.as_str()),
            encode(conversation_id),
            options.model,
            options.restart,
        );
        let uri = Uri::try_from(uri.as_str()).unwrap();
        let body = serde_json::to_vec(&serde_json::json!({
            "sentence": sentence,
            "system_prompt": options.system_prompt
        }))
        .unwrap_or_default();
        match Request::new(&uri)
            .method(Method::POST)
            .header("Content-Type", "application/json")
            .header("Content-Length", &body.len())
            .body(&body)
            .send(&mut writer)
        {
            Ok(res) => {
                match res.status_code().is_success() {
                    true => serde_json::from_slice::<ChatResponse>(&writer)
                        .or(Err(String::from("Unexpected error"))),
                    false => {
                        match res.status_code().into() {
                            409 | 429 | 503 => {
                                // 409 TryAgain 429 RateLimitError
                                // 503 ServiceUnavableila
                                if retry_times > 1 {
                                    sleep(Duration::from_secs(RETRY_INTERVAL));
                                    return chat_completion_inner(
                                        account,
                                        conversation_id,
                                        sentence,
                                        options,
                                        retry_times - 1,
                                    );
                                }
                            }
                            _ => {}
                        }
                        Err(String::from_utf8_lossy(&writer).into_owned())
                    }
                }
            }
            Err(e) => Err(e.to_string()),
        }
    }
}

#[derive(Debug)]
pub enum ImageSize {
    S256,
    S512,
    S1024,
}

impl fmt::Display for ImageSize {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            ImageSize::S256 => write!(f, "256x256"),
            ImageSize::S512 => write!(f, "512x512"),
            ImageSize::S1024 => write!(f, "1024x1024"),
        }
    }
}

impl Serialize for ImageSize {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        match self {
            ImageSize::S256 => serializer.serialize_str("256x256"),
            ImageSize::S512 => serializer.serialize_str("512x512"),
            ImageSize::S1024 => serializer.serialize_str("1024x1024"),
        }
    }
}

/// Request struct for the image creation.
///
/// Use retry_times to set the number of retries when requesting
/// OpenAI's api encounters a problem. The max number is 10.
///
/// For more detail about parameters, please refer to
/// [OpenAI docs](https://platform.openai.com/docs/api-reference/images/create)
///
#[derive(Debug, Serialize)]
pub struct ImageRequest {
    pub prompt: String,
    pub n: u8,
    pub size: ImageSize,
    pub retry_times: u8,
}

/// Create image for the provided prompt and parameters.
///
/// `account` is an [FlowsAccount] used for picking your tied OpenAI API key.
///
/// `params` is a [ImageRequest] object.
///
/// If you have not connected your OpenAI account with [Flows.network platform](https://flows.network),
/// you will receive an error in the flow's building log or running log.
///
pub async fn create_image(
    account: FlowsAccount,
    params: ImageRequest,
) -> Result<Vec<String>, String> {
    let retry_times = match params.retry_times {
        r if r <= 0 => 1,
        r if r > MAX_RETRY_TIMES => MAX_RETRY_TIMES,
        r => r,
    };
    let account = match account {
        FlowsAccount::Default => String::new(),
        FlowsAccount::Provided(s) => s,
    };
    create_image_inner(account, params, retry_times)
}

fn create_image_inner(
    account: String,
    params: ImageRequest,
    retry_times: u8,
) -> Result<Vec<String>, String> {
    unsafe {
        let mut flows_user = Vec::<u8>::with_capacity(100);
        let c = get_flows_user(flows_user.as_mut_ptr());
        flows_user.set_len(c as usize);
        let flows_user = String::from_utf8(flows_user).unwrap();

        let mut writer = Vec::new();
        let uri = format!(
            "{}/{}/create_image?account={}",
            OPENAI_API_PREFIX.as_str(),
            flows_user,
            encode(account.as_str()),
        );
        let uri = Uri::try_from(uri.as_str()).unwrap();
        let body = serde_json::to_vec(&params).unwrap_or_default();
        match Request::new(&uri)
            .method(Method::POST)
            .header("Content-Type", "application/json")
            .header("Content-Length", &body.len())
            .body(&body)
            .send(&mut writer)
        {
            Ok(res) => {
                match res.status_code().is_success() {
                    true => serde_json::from_slice::<Vec<String>>(&writer)
                        .or(Err(String::from("Unexpected error"))),
                    false => {
                        match res.status_code().into() {
                            409 | 429 | 503 => {
                                // 409 TryAgain 429 RateLimitError
                                // 503 ServiceUnavableila
                                if retry_times > 1 {
                                    sleep(Duration::from_secs(RETRY_INTERVAL));
                                    return create_image_inner(account, params, retry_times - 1);
                                }
                            }
                            _ => {}
                        }
                        Err(String::from_utf8_lossy(&writer).into_owned())
                    }
                }
            }
            Err(e) => Err(e.to_string()),
        }
    }
}

/// The input type for the embeddings.
///
#[derive(Debug, Serialize)]
pub enum EmbeddingsInput {
    String(String),
    Vec(Vec<String>),
}

/// Request struct for the embeddings.
///
/// Use retry_times to set the number of retries when requesting
/// OpenAI's api encounters a problem. The max number is 10.
/// For more detail about parameters, please refer to
/// [OpenAI docs](https://platform.openai.com/docs/api-reference/embeddings/create)
///
#[derive(Debug, Serialize)]
pub struct EmbeddingsRequest {
    pub input: EmbeddingsInput,
    #[serde(skip_serializing)]
    pub retry_times: u8,
}

/// Create embeddings from the provided input.
///
/// `account` is an [FlowsAccount] used for picking your tied OpenAI API key.
///
/// `params` is a [EmbeddingsRequest] object.
///
/// If you have not connected your OpenAI account with [Flows.network platform](https://flows.network),
/// you will receive an error in the flow's building log or running log.
///
pub async fn create_embeddings(
    account: FlowsAccount,
    params: EmbeddingsRequest,
) -> Result<Vec<Vec<f64>>, String> {
    let retry_times = match params.retry_times {
        r if r <= 0 => 1,
        r if r > MAX_RETRY_TIMES => MAX_RETRY_TIMES,
        r => r,
    };
    let account = match account {
        FlowsAccount::Default => String::new(),
        FlowsAccount::Provided(s) => s,
    };
    create_embeddings_inner(account, params, retry_times)
}

fn create_embeddings_inner(
    account: String,
    params: EmbeddingsRequest,
    retry_times: u8,
) -> Result<Vec<Vec<f64>>, String> {
    unsafe {
        let mut flows_user = Vec::<u8>::with_capacity(100);
        let c = get_flows_user(flows_user.as_mut_ptr());
        flows_user.set_len(c as usize);
        let flows_user = String::from_utf8(flows_user).unwrap();

        let mut writer = Vec::new();
        let uri = format!(
            "{}/{}/create_embeddings?account={}",
            OPENAI_API_PREFIX.as_str(),
            flows_user,
            encode(account.as_str()),
        );
        let uri = Uri::try_from(uri.as_str()).unwrap();
        let body = match params.input {
            EmbeddingsInput::String(ref s) => serde_json::to_vec(&s).unwrap_or_default(),
            EmbeddingsInput::Vec(ref v) => serde_json::to_vec(&v).unwrap_or_default(),
        };
        match Request::new(&uri)
            .method(Method::POST)
            .header("Content-Type", "application/json")
            .header("Content-Length", &body.len())
            .body(&body)
            .send(&mut writer)
        {
            Ok(res) => {
                match res.status_code().is_success() {
                    true => serde_json::from_slice::<Vec<Vec<f64>>>(&writer)
                        .or(Err(String::from("Unexpected error"))),
                    false => {
                        match res.status_code().into() {
                            409 | 429 | 503 => {
                                // 409 TryAgain 429 RateLimitError
                                // 503 ServiceUnavableila
                                if retry_times > 1 {
                                    sleep(Duration::from_secs(RETRY_INTERVAL));
                                    return create_embeddings_inner(
                                        account,
                                        params,
                                        retry_times - 1,
                                    );
                                }
                            }
                            _ => {}
                        }
                        Err(String::from_utf8_lossy(&writer).into_owned())
                    }
                }
            }
            Err(e) => Err(e.to_string()),
        }
    }
}