//! 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::{CompletionRequest, create_completion};
//! use slack_flows::{listen_to_channel, send_message_to_channel};
//!
//! #[no_mangle]
//! pub fn run() {
//!     listen_to_channel("myworkspace", "mychannel", |sm| {
//!         let cr = CompletionRequest {
//!             prompt: sm.text,
//!             ..Default::default()
//!         };
//!         let r = create_completion("myaccount", cr);
//!         r.iter().for_each(|c| {
//!             send_message_to_channel("myworkspace", "mychannel", c.to_string());
//!         });
//!     });
//! }
//! ```
//!
//! When the Slack message is received, create completion
//! using [create_completion] then send the response to Slack.

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-integration.vercel.app/api")
    );
}

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

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

/// 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,
    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 the account name when you connect
/// [Flows.network](https://flows.network) platform with your OpenAI account.
///
/// `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 fn create_completion(account: &str, params: CompletionRequest) -> Vec<String> {
    let retry_times = match params.retry_times {
        r if r <= 0 => 1,
        r if r > MAX_RETRY_TIMES => MAX_RETRY_TIMES,
        r => r,
    };
    create_completion_inner(account, params, retry_times)
}

pub fn create_completion_inner(
    account: &str,
    params: CompletionRequest,
    retry_times: u8,
) -> Vec<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),
        );
        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) => {
                if !res.status_code().is_success() {
                    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);
                            }
                        }
                        _ => {}
                    }
                    set_error_log(writer.as_ptr(), writer.len() as i32);
                }
                serde_json::from_slice(&writer).unwrap_or_default()
            }
            Err(_) => {
                vec![]
            }
        }
    }
}

/// 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.
///
/// `conversation_id` is the identity of the conversation. The history will be fetched and attached
/// to the `sentence` as a whole prompt for ChatGPT.
///
/// `account` is the account name when you connect
/// [Flows.network](https://flows.network) platform with your OpenAI account.
///
/// `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 fn chat_completion(
    account: &str,
    conversation_id: &str,
    sentence: &str,
    options: &ChatOptions,
) -> Option<ChatResponse> {
    let retry_times = match options.retry_times {
        r if r <= 0 => 1,
        r if r > MAX_RETRY_TIMES => MAX_RETRY_TIMES,
        r => r,
    };
    chat_completion_inner(account, conversation_id, sentence, options, retry_times)
}

pub fn chat_completion_inner(
    account: &str,
    conversation_id: &str,
    sentence: &str,
    options: &ChatOptions,
    retry_times: u8,
) -> Option<ChatResponse> {
    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),
            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) => {
                if !res.status_code().is_success() {
                    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,
                                );
                            }
                        }
                        _ => {}
                    }
                    set_error_log(writer.as_ptr(), writer.len() as i32);
                }
                serde_json::from_slice::<ChatResponse>(&writer).ok()
            }
            Err(_) => None,
        }
    }
}

#[derive(Debug, Serialize)]
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"),
        }
    }
}

/// Request struct for the image creation.
///
/// 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,
}

/// Create image for the provided prompt and parameters.
///
/// `account` is the account name when you connect
/// [Flows.network](https://flows.network) platform with your OpenAI account.
///
/// `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 fn create_image(account: &str, params: ImageRequest) -> Vec<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),
        );
        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) => {
                if !res.status_code().is_success() {
                    set_error_log(writer.as_ptr(), writer.len() as i32);
                }
                serde_json::from_slice(&writer).unwrap_or_default()
            }
            Err(_) => {
                vec![]
            }
        }
    }
}