#[derive(Debug, PartialEq)]
pub enum CallbackReason {
    Start,
    Data,
    End,
    Error(String),
}

#[derive(Debug, PartialEq)]
pub enum Error {
    CurlError(curl::Error),
    FromUtf8Error(std::string::FromUtf8Error),
    SerdeJsonError(String),
}

impl From<curl::Error> for Error {
    fn from(err: curl::Error) -> Self {
        Self::CurlError(err)
    }
}

impl From<std::string::FromUtf8Error> for Error {
    fn from(err: std::string::FromUtf8Error) -> Self {
        Self::FromUtf8Error(err)
    }
}

impl From<serde_json::Error> for Error {
    fn from(err: serde_json::Error) -> Self {
        Self::SerdeJsonError(err.to_string())
    }
}

#[derive(Debug, serde::Deserialize)]
pub struct Permission {
    pub id: String,
    pub object: String,
    pub created: u64,
    pub allow_create_engine: bool,
    pub allow_sampling: bool,
    pub allow_logprobs: bool,
    pub allow_search_indices: bool,
    pub allow_view: bool,
    pub allow_fine_tuning: bool,
    pub organization: String,
    pub group: serde_json::Value,
    pub is_blocking: bool,
}

#[derive(Debug, serde::Deserialize)]
pub struct Model {
    pub id: String,
    pub object: String,
    pub created: u64,
    pub owned_by: String,
    pub permission: Vec<Permission>,
    pub root: String,
    pub parent: Option<String>,
}

#[derive(Debug, serde::Deserialize)]
pub struct ModelList {
    pub object: String,
    pub data: Vec<Model>,
}

#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
pub struct Message {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub role: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub content: Option<String>,
}

#[derive(Debug, Clone, serde::Serialize)]
pub struct RequestBody {
    pub model: String,
    pub messages: Vec<Message>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stream: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub user: Option<String>,
}

#[derive(Debug, serde::Deserialize, serde::Serialize)]
pub struct Choice {
    pub index: u64,
    pub delta: Message,
    pub finish_reason: Option<String>,
}

#[derive(Debug, serde::Deserialize, serde::Serialize)]
pub struct Completion {
    pub id: String,
    pub object: String,
    pub created: u64,
    pub model: String,
    pub choices: Vec<Choice>,
}

mod internal {
    use curl;

    pub fn init(api_key: &str, url: &str) -> Result<curl::easy::Easy, super::Error> {
        let mut easy = curl::easy::Easy::new();
        easy.url(url)?;
        let mut headers = curl::easy::List::new();
        headers.append(&format!("Authorization: Bearer {}", api_key))?;
        headers.append("Content-Type: application/json")?;
        headers.append("Accept: text/event-stream")?;
        easy.http_headers(headers)?;
        Ok(easy)
    }
}

pub mod ll {
    pub async fn list_models<F>(api_key: &str, f: F) -> Result<(), super::Error>
    where
        F: Fn(&[u8]),
    {
        let mut easy = super::internal::init(api_key, "https://api.openai.com/v1/models")?;
        easy.get(true)?;
        let mut transfer = easy.transfer();
        transfer.write_function(|data| {
            f(data);
            Ok(data.len())
        })?;
        transfer.perform()?;
        Ok(())
    }

    pub async fn completions<F>(
        api_key: &str,
        request_body: &super::RequestBody,
        f: F,
    ) -> Result<(), super::Error>
    where
        F: Fn(&[u8]),
    {
        let string_body = serde_json::to_string(request_body)?;
        let mut easy =
            super::internal::init(api_key, "https://api.openai.com/v1/chat/completions")?;
        easy.post(true)?;
        easy.post_fields_copy(string_body.as_bytes())?;
        let mut transfer = easy.transfer();
        transfer.write_function(|data| {
            f(data);
            Ok(data.len())
        })?;
        transfer.perform()?;
        Ok(())
    }
}

pub mod hl {
    pub async fn list_models(api_key: &str) -> Result<String, super::Error> {
        let amv = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
        let cloned_amv = amv.clone();
        super::ll::list_models(&api_key, move |data| {
            let mut v = cloned_amv.lock().unwrap();
            v.extend_from_slice(data)
        })
        .await?;
        let v = amv.lock().unwrap().clone();
        Ok(String::from_utf8(v)?)
    }

    pub async fn completions<F>(
        api_key: &str,
        request_body: &super::RequestBody,
        f: F,
    ) -> Result<(), super::Error>
    where
        F: Fn(String),
    {
        super::ll::completions(&api_key, request_body, |data| {
            f(String::from_utf8(data.to_vec()).unwrap());
        })
        .await?;
        Ok(())
    }
}

pub async fn list_models(api_key: &str) -> Result<Vec<ModelList>, Error> {
    let json = hl::list_models(api_key).await?;
    Ok(serde_json::from_str(&json)?)
}

pub async fn completions<F>(api_key: &str, request_body: &RequestBody, f: F) -> Result<(), Error>
where
    F: Fn(CallbackReason, Option<Completion>),
{
    hl::completions(&api_key, request_body, |data| {
        let is_debug = std::env::var("OPENAI_DEBUG").unwrap_or(String::from("")) != "";
        data.lines().for_each(|line| {
            let prefix = "data: {";
            let prefix_len = prefix.len();
            if line.starts_with(prefix) {
                let string_json = match line.char_indices().nth(prefix_len - 1) {
                    Some((i, _)) => &line[i..],
                    None => "",
                };
                if is_debug {
                    eprintln!("{}", string_json);
                }
                let completion: Completion = serde_json::from_str(string_json).unwrap();
                if completion.choices[0].delta.role == Some(String::from("assistant")) {
                    f(CallbackReason::Start, Some(completion));
                } else if completion.choices[0].finish_reason == Some(String::from("stop")) {
                    f(CallbackReason::End, Some(completion));
                } else {
                    f(CallbackReason::Data, Some(completion));
                }
            } else {
                if is_debug {
                    eprintln!("{}", line);
                }
                f(CallbackReason::Error(String::from(line)), None);
            }
        });
    })
    .await?;
    Ok(())
}

#[cfg(test)]
mod tests {
    #[test]
    fn init_handle() {
        let api_key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY is not defined");
        let result = crate::internal::init(&api_key, "https://api.openai.com/v1/models");
        assert!(result.is_ok());
    }

    #[test]
    fn ll_list_models() {
        let api_key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY is not defined");
        let future = crate::ll::list_models(&api_key, |_| {});
        let result = futures::executor::block_on(future);
        assert!(result.is_ok());
    }

    #[test]
    fn hl_list_models() {
        let api_key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY is not defined");
        let future = crate::hl::list_models(&api_key);
        let result = futures::executor::block_on(future);
        assert!(result.is_ok());
    }

    #[test]
    fn ll_completions() {
        let api_key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY is not defined");
        let request_body = super::RequestBody {
            model: String::from("gpt-3.5-turbo"),
            messages: vec![super::Message {
                role: Some(String::from("user")),
                content: Some(String::from("Say hello")),
            }],
            temperature: None,
            stream: Some(true),
            user: None,
        };
        let future = super::ll::completions(&api_key, &request_body, |_| {});
        let result = futures::executor::block_on(future);
        assert!(result.is_ok());
    }

    #[test]
    fn hl_completions() {
        let api_key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY is not defined");
        let request_body = super::RequestBody {
            model: String::from("gpt-3.5-turbo"),
            messages: vec![super::Message {
                role: Some(String::from("user")),
                content: Some(String::from("Say hello")),
            }],
            temperature: None,
            stream: Some(true),
            user: None,
        };
        let count_start = std::cell::Cell::new(0);
        let count_data = std::cell::Cell::new(0);
        let count_end = std::cell::Cell::new(0);
        let future = super::hl::completions(&api_key, &request_body, |data| {
            data.lines().for_each(|line| {
                let prefix = "data: {";
                let prefix_len = prefix.len();
                if line.starts_with(prefix) {
                    let string_json = match line.char_indices().nth(prefix_len - 1) {
                        Some((i, _)) => &line[i..],
                        None => "",
                    };
                    let completion: super::Completion = serde_json::from_str(string_json).unwrap();
                    if completion.choices[0].delta.role == Some(String::from("assistant")) {
                        count_start.set(count_start.get() + 1);
                    } else if completion.choices[0].finish_reason == Some(String::from("stop")) {
                        count_end.set(count_end.get() + 1);
                    } else {
                        count_data.set(count_data.get() + 1);
                    }
                }
            });
        });
        let result = futures::executor::block_on(future);
        assert!(result.is_ok());
        assert_eq!(count_start.get(), 1);
        assert!(count_data.get() > 0);
        assert_eq!(count_end.get(), 1);
    }

    #[test]
    fn completions() {
        let api_key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY is not defined");
        let request_body = super::RequestBody {
            model: String::from("gpt-3.5-turbo"),
            messages: vec![super::Message {
                role: Some(String::from("user")),
                content: Some(String::from("Say hello")),
            }],
            temperature: None,
            stream: Some(true),
            user: None,
        };
        let count_start = std::cell::Cell::new(0);
        let count_data = std::cell::Cell::new(0);
        let count_end = std::cell::Cell::new(0);
        let future = super::completions(&api_key, &request_body, |cr, completion| match cr {
            super::CallbackReason::Start => {
                count_start.set(count_start.get() + 1);
                assert!(completion.is_some());
                assert_eq!(completion.unwrap().choices.len(), 1);
            }
            super::CallbackReason::Data => {
                count_data.set(count_data.get() + 1);
                assert!(completion.is_some());
                assert_eq!(completion.unwrap().choices.len(), 1);
            }
            super::CallbackReason::End => {
                count_end.set(count_end.get() + 1);
                assert!(completion.is_some());
                assert_eq!(completion.unwrap().choices.len(), 1);
            }
            super::CallbackReason::Error(_) => {
                assert!(completion.is_none());
            }
        });
        let result = futures::executor::block_on(future);
        assert!(result.is_ok());
        assert_eq!(count_start.get(), 1);
        assert!(count_data.get() > 0);
        assert_eq!(count_end.get(), 1);
    }
}