chatgpt_api 0.3.0

use std::env;
use std::io::Write;
use std::path::Path;

use chatgpt::chatbot::Chatbot;
use chatgpt::config::Config;
use chatgpt::error::GPTError;
use dialoguer::theme::ColorfulTheme;
use dialoguer::Input;
use futures::StreamExt;

pub mod args;

fn get_input(prompt: &str) -> String {
    Input::<String>::with_theme(&ColorfulTheme::default())
        .with_prompt(prompt)
        .interact_text()
        .expect("Failed to get input")
}

pub fn configure() -> Config {
    let mut config_files = vec!["config.json".to_string()];
    let xdg_config_home = env::var("XDG_CONFIG_HOME").ok();
    if let Some(xdg_config_home) = xdg_config_home {
        config_files.push(format!("{xdg_config_home}/chatgpt/config.json"));
    }
    let user_home = env::var("HOME").ok();
    if let Some(user_home) = user_home {
        config_files.push(format!("{user_home}/.config/chatgpt/config.json"));
    }
    let config_file = config_files.iter().find(|f| Path::new(f).exists());
    config_file.map_or_else(
        || {
            panic!("No config file found.");
        },
        |config_file| Config::from_file(Path::new(config_file)),
    )
}

fn handle_commands(command: &str, chatbot: &mut Chatbot) -> bool {
    if command == "!help" {
        println!(
            "
        !help - Show this message
        !reset - Forget the current conversation
        !config - Show the current configuration
        !rollback x - Rollback the conversation (x being the number of messages to rollback)
        !exit - Exit this program
        !setconversation - Changes the conversation
        "
        );
    }
    else if command == "!reset" {
        chatbot.reset_chat();
        println!("Chat session successfully reset.");
    }
    else if command == "!config" {
        println!("{:#?}", chatbot.config);
    }
    else if command.starts_with("!rollback") {
        // Default to 1 rollback if no number is specified
        let rollback = command.split(' ').nth(1).map_or_else(
            || {
                println!("No number specified, rolling back 1 message");
                1
            },
            |rollback| rollback.parse::<usize>().unwrap(),
        );
        chatbot.rollback_conversation(rollback);
        println!("Rolled back {rollback} messages.");
    }
    else if command.starts_with("!setconversation") {
        match command.split(' ').nth(1) {
            Some(conversation_id) => {
                chatbot.conversation_id = Some(conversation_id.parse().unwrap());

                chatbot.config.conversation_id = Some(conversation_id.to_string());
                println!("Conversation has been changed");
            },
            None => {
                println!("Please include conversation UUID in command");
            },
        }
    }
    else if command == "!exit" {
        std::process::exit(0);
    }
    else {
        return false;
    }
    true
}

pub async fn run(config: Config) -> Result<(), GPTError> {
    log::info!("Logging in...");
    let mut chatbot = Chatbot::new(config.clone()).await;
    log::info!("Logged in! ✅");

    let stdout = std::io::stdout();

    loop {
        let prompt = get_input("You: ");

        // Check if the prompt is a command
        if prompt.starts_with('!') && handle_commands(&prompt, &mut chatbot) {
            continue;
        }

        let mut message_stream = chatbot.ask_stream(&prompt).await;
        let mut last_message = String::new();
        while let Some(m) = message_stream.next().await {
            let message = m?;
            assert!(
                message.starts_with(&last_message),
                "Message doesn't start with last message: {message:?} {last_message:?}"
            );
            print!("{}", message.strip_prefix(&last_message).unwrap());
            stdout.lock().flush().unwrap();
            last_message = message;
        }
        println!();
    }
}