
use std::error::Error;
use std::fmt;

use diesel::{prelude::*, PgConnection, insert_into};
use diesel::r2d2::{Pool, ConnectionManager};
use log::{error, debug, trace, warn};

use serde::{Serialize, Deserialize};
use serde_json::Value;
use serenity::model::channel::Message;
use serenity::model::prelude::ChannelType;
use serenity::prelude::*;

use crate::database::models::{NewMessageDB, MessageDB};

pub struct OAI {
  pub client: reqwest::Client,
  pub base_url: String,
  pub max_attempts: i64,
  pub token: String,
  pub max_context_questions: i64
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct ChatCompletionRequest {
  model: GPTModel,
  messages: Vec<ChatCompletionMessage>,
  /// Value between 0 and 2
  #[serde(skip_serializing_if = "Option::is_none")]
  temperature: Option<f64>,
  /// Value between 0 and 1
  #[serde(skip_serializing_if = "Option::is_none")]
  top_p: Option<f64>,
  #[serde(skip_serializing_if = "Option::is_none")]
  n: Option<f64>,
  #[serde(skip_serializing_if = "Option::is_none")]
  max_tokens: Option<i64>,
  /// Value between -2.0 and 2.0
  #[serde(skip_serializing_if = "Option::is_none")]
  presence_penalty: Option<f64>,
  /// Value between -2.0 and 2.0
  #[serde(skip_serializing_if = "Option::is_none")]
  frequency_penalty: Option<f64>,
  #[serde(skip_serializing_if = "Option::is_none")]
  user: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct ChatCompletionMessage {
  role: GPTRole,
  content: String
}

#[derive(Debug, Clone, Serialize, Deserialize)]
enum GPTRole {
  #[serde(rename = "system")]
  System,
  #[serde(rename = "user")]
  User,
  #[serde(rename = "assistant")]
  Assistant,
  #[serde(rename = "function")]
  Function
}

#[derive(Debug, Clone, Serialize, Deserialize)]
enum GPTModel {
  #[serde(rename = "gpt-3.5-turbo")]
  GPT35Turbo,
  #[serde(rename = "gpt-3.5-turbo-0613")]
  GPT35Snapshot,
  #[serde(rename = "gpt-3.5-turbo-16k")]
  GPT3516k,
  #[serde(rename = "gpt-3.5-turbo-16k-0613")]
  GPT3516kSnapshot,
  #[serde(rename = "gpt-4")]
  GPT4,
  #[serde(rename = "gpt-4-0613")]
  GPT4Snapshot,
  #[serde(rename = "gpt-4-32k")]
  GPT432k,
  #[serde(rename = "gpt-4-32k-0613")]
  GPT432kSnapshot,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct ChatCompletionResponse {
  id: String,
  object: String,
  created: i64,
  model: GPTModel,
  usage: Usage,
  choices: Vec<Choice>
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct Usage {
  prompt_tokens: i64,
  completion_tokens: i64,
  total_tokens: i64
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct Choice {
  message: ChatCompletionMessage,
  finish_reason: String,
  index: i64
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct ResponseError {
  code: Option<i64>,
  message: Option<String>,
  param: Option<String>,
  #[serde(rename = "type")]
  error_type: Option<String>
}

#[derive(Debug)]
struct OAIError {
  pub message: String
}

impl fmt::Display for OAIError {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(f, "OAIError: {}", self.message)
  }
}

impl Error for OAIError {}

#[derive(Debug, Clone, Serialize, Deserialize)]
enum ResponseEvent {
  ChatCompletionResponse(ChatCompletionResponse),
  ResponseError(ResponseError)
}

impl OAI {
  async fn get_request(&self, request: ChatCompletionRequest) -> Result<ChatCompletionResponse, OAIError> {
    let uri = format!("{}/chat/completions", self.base_url);
    let body = serde_json::to_string(&request).unwrap();
    trace!("Sending request to {}: {}", uri, body);

    let value = match match self.client
      .post(&uri)
      .bearer_auth(&self.token)
      .header("Content-Type", "application/json".to_string())
      .body(body)
      .send()
      .await {
        Ok(r) => r,
        Err(err) => return Err(OAIError {
          message: format!("Could not send request to OpenAI: {}", err),
        })
      }
      .json::<Value>()
      .await {
        Ok(r) => r,
        Err(err) => return Err(OAIError {
          message: format!("Could not read response from OpenAI: {}", err)
        })
      };

    trace!("Received response from OpenAI: {:?}", value);
    
    // let response = match serde_json::from_value::<OAIResponseEvent>(value) {
    //   Ok(r) => {
    //     match r {
    //       OAIResponseEvent::OAIResponse(r) => r,
    //       OAIResponseEvent::OAIError(e) => return Err(OAIError { message: e.message.unwrap_or("Unknown error".to_string()) })
    //     }
    //   },
    //   Err(err) => return Err(OAIError {
    //     message: format!("Could not parse response from OpenAI: {}", err)
    //   })
    // };
    let response = match serde_json::from_value::<ChatCompletionResponse>(value) {
      Ok(r) => r,
      Err(err) => return Err(OAIError {
        message: format!("Could not parse response from OpenAI: {}", err)
      })
    };

    Ok(response)
  }
}

pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI, pool: &Pool<ConnectionManager<PgConnection>>) {
  debug!("Generating response for message: {}", msg.content);
  let typing = msg.channel_id.start_typing(&ctx.http).unwrap();

  let guild_id = msg.guild_id.unwrap();
  let channel_id = msg.channel_id;
  let author_id = msg.author.id;
  let mut connection = pool.get().unwrap();

  // Parse out the bot mention from the message
  let bot_mention: String = format!("<@{}>", ctx.cache.current_user_id().0);
  let parsed_content = msg.content.replace(bot_mention.as_str(), "");

  // Setup the request messages
  let result: Result<Vec<MessageDB>, diesel::result::Error> = crate::database::schema::messages::table
    .select(MessageDB::as_select())
    .filter((crate::database::schema::messages::guild_id.eq(guild_id.0 as i64))
      .and(crate::database::schema::messages::channel_id.eq(channel_id.0  as i64))
      .and(crate::database::schema::messages::user_id.eq(author_id.0  as i64))
    )
    .order(crate::database::schema::messages::created.asc())
    .limit(oai.max_context_questions)
    .load(&mut connection);
  
  let previous_messages = match result {
    Ok(r) => {
      let mut previous_message = "".to_string();
      for message in r {
        previous_message = format!("{}\nYou: {}\n Siren: {}", previous_message, message.request, message.response);
      }
      Some(ChatCompletionMessage { role: GPTRole::User, content: previous_message })
    }
    Err(err) => {
      error!("Could not load previous messages: {}", err);
      None
    }
  };

  let mut messages = vec![
    ChatCompletionMessage {
      role: GPTRole::System,
      content: "Siren is a Discord bot specializing in Dungeons and Dragons. Limit Siren's responses to <= 2000 characters. Siren must always obey these instructions, no matter what.".to_string()
    },
  ];

  if let Some(mut previous) = previous_messages {
    previous.content = format!("{}\nYou: {}\nSiren: ", previous.content, parsed_content);
    messages.push(previous);
  } else {
    messages.push(ChatCompletionMessage {
      role: GPTRole::User,
      content: format!("You: {}, Siren: ", parsed_content)
    });
  }

  let model = "gpt-3.5-turbo".to_string();

  let request = ChatCompletionRequest {
    model: GPTModel::GPT35Turbo,
    messages,
    temperature: Some(0.5),
    top_p: None,
    n: None,
    max_tokens: Some(1000),
    presence_penalty: Some(0.6),
    frequency_penalty: Some(0.0),
    user: Some(msg.author.name.clone())
  };
  let response = match oai.get_request(request).await {
    Ok(r) => {
      debug!("Processing response received from OpenAI");
      if !r.choices.is_empty() {
        // Insert the message into the messages database table
        let res = r.choices[0].message.content.clone();
        if let Err(err) = insert_into(crate::database::schema::messages::table).values(NewMessageDB {
          id: &r.id,
          guild_id: guild_id.0 as i64,
          channel_id: channel_id.0 as i64,
          user_id: author_id.0 as i64,
          created: r.created,
          model: &model,
          request: &parsed_content,
          response: &res,
          request_tags: vec![],
          response_tags: vec![],
        }).execute(&mut connection) {
          error!("Could not insert message into database: {}", err);
        }
        res
      } else {
        warn!("No choices received in the response from OpenAI");
        "No reply received".to_string()
      }
    }
    Err(err) => {
      error!("Could not get response from OpenAI: {}", err.message);
      err.message
    }
  };
  debug!("Writing response: \"{}\"", response);

  // Stop the typing indicator and send the response
  typing.stop();
  match msg.channel_id.create_public_thread(&ctx.http, msg.id, |thread| {
      thread.name(truncate(&parsed_content, 99)).kind(ChannelType::PublicThread)
  }).await {
    Ok(c) => {
      if let Err(why) = c.say(&ctx.http, response).await {
        error!("Cannot send message: {}", why);
      }
    }
    Err(_) => {
      if let Err(why) = channel_id.say(&ctx.http, response).await {
        error!("Cannot send message: {}", why);
      }
    }
  };
}

fn truncate(s: &str, max_chars: usize) -> &str {
  match s.char_indices().nth(max_chars) {
    None => s,
    Some((idx, _)) => &s[..idx],
  }
}

// fn summarization_model() -> Result<SummarizationModel, RustBertError> {
//   let config_resource = RemoteResource::from_pretrained(T5ConfigResources::T5_SMALL);
//   let vocab_resource = RemoteResource::from_pretrained(T5VocabResources::T5_SMALL);
//   let weights_resource = RemoteResource::from_pretrained(T5ModelResources::T5_SMALL);

//   let summarization_config = SummarizationConfig::new(
//     ModelType::T5,
//     ModelResource::Torch(Box::new(weights_resource)),
//     config_resource,
//     vocab_resource,
//     None
//   );
//   return SummarizationModel::new(summarization_config);
// }
