rustGPT/src/cli.rs

use anyhow::Result;

use crate::config::Config;
use crate::core::{
    create_client, get_provider_for_model, provider::get_all_models, provider::get_supported_models,
    provider::is_model_supported, ChatClient, Session,
};
use crate::utils::{Display, InputHandler};

pub struct ChatCLI {
    session: Session,
    client: Option<ChatClient>,
    current_model: Option<String>,
    display: Display,
    input: InputHandler,
    config: Config,
}

impl ChatCLI {
    pub fn new(session: Session, config: Config) -> Result<Self> {
        Ok(Self {
            session,
            client: None,
            current_model: None,
            display: Display::new(),
            input: InputHandler::new()?,
            config,
        })
    }

    fn get_client(&mut self) -> Result<&ChatClient> {
        if self.client.is_none() || self.current_model.as_ref() != Some(&self.session.model) {
            let client = create_client(&self.session.model, &self.config)?;
            self.current_model = Some(self.session.model.clone());
            self.client = Some(client);
        }
        Ok(self.client.as_ref().unwrap())
    }

    pub async fn run(&mut self) -> Result<()> {
        self.display.print_header();
        self.display.print_info("Type your message and press Enter. Commands start with '/'.");
        self.display.print_info("Type /help for help.");
        
        let provider = get_provider_for_model(&self.session.model);
        self.display.print_model_info(&self.session.model, provider.as_str());
        self.display.print_session_info(&self.session.name);
        
        println!();

        loop {
            match self.input.read_line("👤> ")? {
                Some(line) => {
                    let line = line.trim();
                    if line.is_empty() {
                        continue;
                    }

                    if line.starts_with('/') {
                        if !self.handle_command(line).await? {
                            break;
                        }
                    } else {
                        if let Err(e) = self.handle_user_message(line).await {
                            self.display.print_error(&format!("Error: {}", e));
                        }
                    }
                }
                None => {
                    self.display.print_info("Goodbye!");
                    break;
                }
            }
        }

        self.session.save()?;
        self.input.save_history()?;
        Ok(())
    }

    async fn handle_user_message(&mut self, message: &str) -> Result<()> {
        self.session.add_user_message(message.to_string());
        self.session.save()?;

        let spinner = self.display.show_spinner("Thinking");
        
        // Clone data needed for the API call before getting mutable client reference
        let model = self.session.model.clone();
        let messages = self.session.messages.clone();
        let enable_web_search = self.session.enable_web_search;
        let enable_reasoning_summary = self.session.enable_reasoning_summary;
        let reasoning_effort = self.session.reasoning_effort.clone();
        
        let client = self.get_client()?;
        
        match client
            .chat_completion(
                &model,
                &messages,
                enable_web_search,
                enable_reasoning_summary,
                &reasoning_effort,
            )
            .await
        {
            Ok(response) => {
                spinner.finish("Done");
                self.display.print_assistant_response(&response);
                self.session.add_assistant_message(response);
                self.session.save()?;
            }
            Err(e) => {
                spinner.finish_with_error("Failed");
                return Err(e);
            }
        }

        Ok(())
    }

    async fn handle_command(&mut self, command: &str) -> Result<bool> {
        let parts: Vec<&str> = command.split_whitespace().collect();
        if parts.is_empty() {
            return Ok(true);
        }

        match parts[0].to_lowercase().as_str() {
            "/help" => {
                self.display.print_help();
            }
            "/exit" => {
                self.session.save()?;
                self.display.print_info("Session saved. Goodbye!");
                return Ok(false);
            }
            "/model" => {
                self.handle_model_command(&parts).await?;
            }
            "/models" => {
                self.list_models();
            }
            "/list" => {
                self.list_sessions()?;
            }
            "/new" => {
                self.handle_new_session(&parts)?;
            }
            "/switch" => {
                self.handle_switch_session(&parts).await?;
            }
            "/clear" => {
                self.session.clear_messages();
                self.session.save()?;
                self.display.print_command_result("Conversation cleared");
            }
            "/delete" => {
                self.handle_delete_session(&parts).await?;
            }
            "/tool" => {
                self.handle_tool_command(&parts)?;
            }
            "/reasoning" => {
                self.handle_reasoning_command(&parts)?;
            }
            "/effort" => {
                self.handle_effort_command(&parts)?;
            }
            _ => {
                self.display.print_error(&format!("Unknown command: {} (see /help)", parts[0]));
            }
        }

        Ok(true)
    }

    async fn handle_model_command(&mut self, parts: &[&str]) -> Result<()> {
        if parts.len() == 1 {
            let all_models = get_all_models();
            let selection = self.input.select_from_list(
                "Select a model:",
                &all_models,
                Some(&self.session.model),
            )?;
            
            if let Some(model) = selection {
                self.session.model = model.to_string();
                let provider = get_provider_for_model(&self.session.model);
                self.display.print_command_result(&format!(
                    "Model switched to {} ({})",
                    self.session.model,
                    provider.as_str()
                ));
                self.client = None; // Force client recreation
            }
        } else if parts.len() == 2 {
            let model = parts[1];
            if !is_model_supported(model) {
                self.display.print_error("Unsupported model. Use /models to see the list of supported models.");
            } else {
                self.session.model = model.to_string();
                let provider = get_provider_for_model(&self.session.model);
                self.display.print_command_result(&format!(
                    "Model switched to {} ({})",
                    self.session.model,
                    provider.as_str()
                ));
                self.client = None; // Force client recreation
            }
        } else {
            self.display.print_error("Usage: /model [model_name]");
        }
        Ok(())
    }

    fn list_models(&self) {
        self.display.print_info("Supported models:");
        let supported = get_supported_models();
        
        for (provider, models) in supported {
            println!("  {}:", provider.as_str().to_uppercase());
            for model in models {
                let marker = if model == self.session.model { " <- current" } else { "" };
                println!("    {}{}", model, marker);
            }
        }
    }

    fn list_sessions(&self) -> Result<()> {
        let sessions = Session::list_sessions()?;
        
        if sessions.is_empty() {
            self.display.print_info("No saved sessions");
            return Ok(());
        }

        self.display.print_info("Saved sessions:");
        for (name, updated) in sessions {
            let marker = if name == self.session.name { "★" } else { " " };
            let date_str = updated.format("%Y-%m-%d %H:%M:%S");
            println!("  {} {} (updated: {})", marker, name, date_str);
        }
        
        Ok(())
    }

    fn handle_new_session(&mut self, parts: &[&str]) -> Result<()> {
        if parts.len() != 2 {
            self.display.print_error("Usage: /new <session_name>");
            return Ok(());
        }

        self.session.save()?;
        let new_session = Session::new(parts[1].to_string(), self.session.model.clone());
        self.session = new_session;
        self.display.print_command_result(&format!("New session '{}' started", self.session.name));
        
        Ok(())
    }

    async fn handle_switch_session(&mut self, parts: &[&str]) -> Result<()> {
        if parts.len() == 1 {
            let sessions = Session::list_sessions()?;
            let session_names: Vec<String> = sessions
                .into_iter()
                .map(|(name, _)| name)
                .filter(|name| name != &self.session.name)
                .collect();

            if let Some(selection) = self.input.select_from_list(
                "Switch to session:",
                &session_names,
                None,
            )? {
                self.session.save()?;
                match Session::load(&selection) {
                    Ok(session) => {
                        self.session = session;
                        self.display.print_command_result(&format!(
                            "Switched to session '{}' (model={})",
                            self.session.name, self.session.model
                        ));
                        self.client = None; // Force client recreation
                    }
                    Err(e) => {
                        self.display.print_error(&format!("Failed to load session: {}", e));
                    }
                }
            }
        } else if parts.len() == 2 {
            let session_name = parts[1];
            self.session.save()?;
            match Session::load(session_name) {
                Ok(session) => {
                    self.session = session;
                    self.display.print_command_result(&format!(
                        "Switched to session '{}' (model={})",
                        self.session.name, self.session.model
                    ));
                    self.client = None; // Force client recreation
                }
                Err(e) => {
                    self.display.print_error(&format!("Failed to load session: {}", e));
                }
            }
        } else {
            self.display.print_error("Usage: /switch [session_name]");
        }
        
        Ok(())
    }

    async fn handle_delete_session(&mut self, parts: &[&str]) -> Result<()> {
        let target = if parts.len() == 1 {
            let sessions = Session::list_sessions()?;
            let session_names: Vec<String> = sessions
                .into_iter()
                .map(|(name, _)| name)
                .filter(|name| name != &self.session.name)
                .collect();

            self.input.select_from_list("Delete session:", &session_names, None)?
        } else if parts.len() == 2 {
            Some(parts[1].to_string())
        } else {
            self.display.print_error("Usage: /delete [session_name]");
            return Ok(());
        };

        if let Some(target) = target {
            if target == self.session.name {
                self.display.print_error(
                    "Cannot delete the session you are currently using. Switch to another session first."
                );
                return Ok(());
            }

            if self.input.confirm(&format!("Delete session '{}'?", target))? {
                match Session::delete_session(&target) {
                    Ok(()) => {
                        self.display.print_command_result(&format!("Session '{}' deleted", target));
                    }
                    Err(e) => {
                        self.display.print_error(&format!("Failed to delete session: {}", e));
                    }
                }
            }
        }
        
        Ok(())
    }

    fn handle_tool_command(&mut self, parts: &[&str]) -> Result<()> {
        if parts.len() != 3 || parts[1].to_lowercase() != "websearch" || !["on", "off"].contains(&parts[2]) {
            self.display.print_error("Usage: /tool websearch on|off");
            return Ok(());
        }

        let enable = parts[2] == "on";
        
        if enable {
            let model = self.session.model.clone();
            if let Ok(client) = self.get_client() {
                if !client.supports_feature_for_model("web_search", &model) {
                    let provider = get_provider_for_model(&model);
                    self.display.print_warning(&format!(
                        "Web search is not supported by {} models",
                        provider.as_str()
                    ));
                }
            }
        }

        self.session.enable_web_search = enable;
        let state = if enable { "enabled" } else { "disabled" };
        self.display.print_command_result(&format!("Web search tool {}", state));
        
        Ok(())
    }

    fn handle_reasoning_command(&mut self, parts: &[&str]) -> Result<()> {
        if parts.len() != 2 || !["on", "off"].contains(&parts[1]) {
            self.display.print_error("Usage: /reasoning on|off");
            return Ok(());
        }

        let enable = parts[1] == "on";
        
        if enable {
            let model = self.session.model.clone();
            if let Ok(client) = self.get_client() {
                if !client.supports_feature_for_model("reasoning_summary", &model) {
                    let provider = get_provider_for_model(&model);
                    self.display.print_warning(&format!(
                        "Reasoning summaries are not supported by {} models",
                        provider.as_str()
                    ));
                }
            }
        }

        self.session.enable_reasoning_summary = enable;
        let state = if enable { "enabled" } else { "disabled" };
        self.display.print_command_result(&format!("Reasoning summaries {}", state));
        
        Ok(())
    }

    fn handle_effort_command(&mut self, parts: &[&str]) -> Result<()> {
        if parts.len() == 1 {
            self.display.print_info(&format!("Current reasoning effort: {}", self.session.reasoning_effort));
            self.display.print_info("Available levels: low, medium, high");
            if !self.session.model.starts_with("gpt-5") {
                self.display.print_warning("Reasoning effort is only supported by GPT-5 models");
            }
        } else if parts.len() == 2 {
            let effort = parts[1];
            if !["low", "medium", "high"].contains(&effort) {
                self.display.print_error("Invalid effort level. Use: low, medium, or high");
            } else {
                if !self.session.model.starts_with("gpt-5") {
                    self.display.print_warning("Reasoning effort is only supported by GPT-5 models");
                }
                self.session.reasoning_effort = effort.to_string();
                self.display.print_command_result(&format!("Reasoning effort set to {}", effort));
            }
        } else {
            self.display.print_error("Usage: /effort [low|medium|high]");
        }
        
        Ok(())
    }
}