tofu/

tofu.rs

//! Tofu - A command-line tool for interacting with LLMs
//!
//! This library provides the core functionality for the Tofu CLI tool.

#![forbid(unsafe_code)]
#![warn(missing_docs)]

use colored::Colorize;
use dialoguer::Editor;
use dialoguer::Input;
use home::home_dir;
use indicatif::{ProgressBar, ProgressStyle};
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::error::Error;
use std::fs;
use std::io::Write;
use std::path::PathBuf;

mod theme;
use theme::TofuTheme;

/// Configuration loaded from config file
#[derive(Debug, Serialize, Deserialize)]
pub struct ConfigFile {
    /// The LLM provider to use
    pub provider: String,
    /// The model to use
    pub model: String,
    /// Whether to JSON stream the response
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stream: Option<bool>,
    /// The system prompt to use
    #[serde(skip_serializing_if = "Option::is_none")]
    pub system_prompt: Option<String>,
}

/// Message in the conversation history
#[derive(Debug, Clone, Serialize, Deserialize)]
struct Message {
    role: String,
    content: String,
}

/// Configuration loaded from keys file
#[derive(Debug, Serialize, Deserialize)]
pub struct KeysFile {
    /// The Google API key
    pub google: Option<String>,
    /// The OpenAI API key
    pub openai: Option<String>,
    /// The Anthropic API key
    pub anthropic: Option<String>,
}

/// The main configuration for the Tofu application.
#[derive(Debug)]
pub struct Config {
    /// Whether to enable verbose output
    pub verbose: bool,
    /// Whether to enable interactive mode
    pub interactive: bool,
    /// Optional message string
    pub message: Option<String>,
    /// Whether to JSON stream the response
    pub stream: Option<bool>,
    /// Configuration loaded from file
    pub file: Option<ConfigFile>,
}

/// Runs the Tofu application with the given configuration.
/// # Arguments
/// * `config` - The configuration for the application
/// # Returns
/// Returns `Ok(())` on success, or an error if something went wrong.
/// Gets the path to the config file.
fn get_config_path() -> Result<PathBuf, Box<dyn Error>> {
    let config_dir = if cfg!(windows) {
        dirs::config_dir()
            .ok_or("Could not determine config directory")?
            .join("tofu")
    } else {
        home_dir()
            .ok_or("Could not determine home directory")?
            .join(".tofu")
    };

    // Create config directory if it doesn't exist
    if !config_dir.exists() {
        std::fs::create_dir_all(&config_dir)?;
    }

    Ok(config_dir.join("config.json"))
}

fn get_keys_path() -> Result<PathBuf, Box<dyn Error>> {
    let config_dir = if cfg!(windows) {
        dirs::config_dir()
            .ok_or("Could not determine config directory")?
            .join("tofu")
    } else {
        home_dir()
            .ok_or("Could not determine home directory")?
            .join(".tofu")
    };

    // Create config directory if it doesn't exist
    if !config_dir.exists() {
        std::fs::create_dir_all(&config_dir)?;
    }

    Ok(config_dir.join("keys.json"))
}

/// Loads the configuration from file.
pub fn load_config(profile: Option<&str>) -> Result<ConfigFile, Box<dyn Error>> {
    let config_path = get_config_path()?;

    if !config_path.exists() {
        // If config file doesn't exist, create it with default values as a multi-profile map
        let default_config = ConfigFile {
            provider: String::from("pollinations"),
            model: String::from("openai"),
            stream: Some(true),
            system_prompt: Some(String::from("You are a helpful assistant named Tofu.")),
        };
        let gemini_config = ConfigFile {
            provider: String::from("google"),
            model: String::from("gemini-2.5-flash"),
            stream: None,
            system_prompt: None,
        };
        let openai_config = ConfigFile {
            provider: String::from("openai"),
            model: String::from("gpt-5-mini"),
            stream: None,
            system_prompt: None,
        };
        let profiles_json = serde_json::json!({
            "default": &default_config,
            "gemini": &gemini_config,
            "openai": &openai_config
        });
        let config_json = serde_json::to_string_pretty(&profiles_json)?;
        std::fs::write(&config_path, config_json)?;
        return Ok(default_config);
    }

    let config_content = fs::read_to_string(&config_path)?;

    // Try to parse as either legacy single-profile or multi-profile config
    let root_value: serde_json::Value = serde_json::from_str(&config_content)
        .map_err(|e| format!("Failed to parse config file: {}", e))?;

    if let Some(obj) = root_value.as_object() {
        let looks_like_legacy = obj.contains_key("provider")
            || obj.contains_key("model")
            || obj.contains_key("stream")
            || obj.contains_key("system_prompt");

        if looks_like_legacy {
            // Legacy single-profile config
            let cfg: ConfigFile = serde_json::from_value(root_value)
                .map_err(|e| format!("Failed to parse legacy config: {}", e))?;
            if cfg.provider.is_empty() || cfg.model.is_empty() {
                return Err("Invalid config: provider and model must not be empty".into());
            }
            return Ok(cfg);
        }

        // Multi-profile config
        let (selected_name, selected_value) = if let Some(name) = profile {
            match obj.get(name) {
                Some(v) => (name.to_string(), v.clone()),
                None => {
                    let available = obj.keys().cloned().collect::<Vec<_>>().join(", ");
                    return Err(
                        format!("Profile '{}' not found. Available: {}", name, available).into(),
                    );
                }
            }
        } else {
            if let Some(v) = obj.get("default") {
                (String::from("default"), v.clone())
            } else {
                match obj.iter().next() {
                    Some((k, v)) => (k.clone(), v.clone()),
                    None => return Err("Config file contains no profiles".into()),
                }
            }
        };

        let mut cfg: ConfigFile = serde_json::from_value(selected_value).map_err(|e| {
            format!(
                "Failed to parse selected profile '{}': {}",
                selected_name, e
            )
        })?;

        // If selected profile is not the default, fall back to default for None values
        if selected_name != "default" {
            if let Some(default_value) = obj.get("default") {
                let default_cfg: ConfigFile = serde_json::from_value(default_value.clone())
                    .map_err(|e| format!("Failed to parse default profile: {}", e))?;
                // Fall back to default for None values
                if cfg.stream.is_none() {
                    cfg.stream = default_cfg.stream;
                }
                if cfg.system_prompt.is_none() {
                    cfg.system_prompt = default_cfg.system_prompt;
                }
            }
        }

        if cfg.provider.is_empty() || cfg.model.is_empty() {
            return Err("Invalid config: provider and model must not be empty".into());
        }
        return Ok(cfg);
    }

    Err("Invalid config: root must be a JSON object".into())
}

/// Loads the keys file from the default location.
pub fn load_keys() -> Result<KeysFile, Box<dyn Error>> {
    let keys_path = get_keys_path()?;

    if !keys_path.exists() {
        // If keys file doesn't exist, create it with default values as a multi-profile map
        let default_keys = serde_json::json!({
            "google": "",
            "openai": "",
            "anthropic": ""
        });

        let keys_json = serde_json::to_string_pretty(&default_keys)?;
        std::fs::write(&keys_path, keys_json)?;
        return Ok(KeysFile {
            google: None,
            openai: None,
            anthropic: None,
        });
    }

    let keys_content = fs::read_to_string(&keys_path)?;
    let keys_json: serde_json::Value = serde_json::from_str(&keys_content)?;
    let keys: KeysFile = serde_json::from_value(keys_json)?;

    return Ok(keys);
}

/// Opens the config file in the default editor.
///
/// # Returns
/// Returns `Ok(())` on success, or an error if something went wrong.
pub fn open_config() -> Result<(), Box<dyn Error>> {
    println!("Opening config file...");
    let config_path = get_config_path()?;

    // Ensure config file exists by trying to load it or create a default one
    if let Err(e) = load_config(None) {
        eprintln!("Warning: {}", e);
        eprintln!("Opening editor to fix the config file...");
    }

    // Open the config file in the default editor
    let editor = std::env::var("EDITOR").unwrap_or_else(|_| {
        if cfg!(windows) {
            String::from("notepad")
        } else {
            String::from("nano")
        }
    });

    let status = std::process::Command::new(editor)
        .arg(&config_path)
        .status()?;

    if !status.success() {
        return Err(format!("Editor exited with status: {}", status).into());
    }

    // Try to load the config after editing, but don't fail if it's still invalid
    if let Err(e) = load_config(None) {
        eprintln!("Warning: The config file is still invalid: {}", e);
        eprintln!("Please fix the config file and try again.");
    }

    Ok(())
}

/// Opens the keys file in the default editor.
pub fn open_keys() -> Result<(), Box<dyn Error>> {
    println!("Opening keys file...");
    let config_path = get_keys_path()?;

    // Ensure keys file exists by trying to load it or create a default one
    if let Err(e) = load_keys() {
        eprintln!("Warning: {}", e);
        eprintln!("Opening editor to fix the keys file...");
    }

    // Open the file in the default editor
    let editor = std::env::var("EDITOR").unwrap_or_else(|_| {
        if cfg!(windows) {
            String::from("notepad")
        } else {
            String::from("nano")
        }
    });

    let status = std::process::Command::new(editor)
        .arg(&config_path)
        .status()?;

    if !status.success() {
        return Err(format!("Editor exited with status: {}", status).into());
    }

    Ok(())
}

/// Runs the Tofu application with the given configuration.
/// # Arguments
/// * `config` - The configuration for the application
/// # Returns
/// Returns `Ok(())` on success, or an error if something went wrong.
/// Runs the Tofu application with the given configuration asynchronously.
/// # Arguments
/// * `config` - The configuration for the application
/// # Returns
/// Returns `Ok(())` on success, or an error if something went wrong.
pub async fn run(config: Config) -> Result<(), Box<dyn Error>> {
    if config.verbose {
        println!(
            "Tofu v{} initialized (verbose mode)",
            env!("CARGO_PKG_VERSION")
        );
        println!("{:#?}", config);
    }

    if config.interactive {
        run_interactive(config).await
    } else {
        let message = config.message.as_ref().unwrap_or(&String::new()).clone();
        send_message(&message, &config, vec![]).await?;
        Ok(())
    }
}

async fn run_interactive(mut config: Config) -> Result<(), Box<dyn Error>> {
    let mut conversation_history = vec![];

    println!(
        "{}",
        format!("Tofu {}", env!("CARGO_PKG_VERSION")).bold().blue()
    );
    println!("{}", "Ctrl+C or /exit to exit".italic().dimmed());

    loop {
        let input: Result<String, _> = Input::with_theme(&TofuTheme::default()).interact_text();

        match input {
            Ok(mut line) => {
                line = line.trim().to_string();
                if line.is_empty() {
                    continue;
                }

                // Check for commands starting with /
                if line.starts_with('/') {
                    let (should_exit, new_config, message_to_send) =
                        handle_command(line.as_str(), &mut conversation_history)?;
                    if let Some(new_file_config) = new_config {
                        config.file = Some(new_file_config);
                    }
                    if should_exit {
                        break; // Exit the loop if command returns true
                    }
                    if let Some(message) = message_to_send {
                        // Process the multiline message like a regular input
                        line = message;
                    } else {
                        continue; // Skip sending to model for commands that don't return a message
                    }
                }

                // Add user message to conversation history
                conversation_history.push(Message {
                    role: "user".to_string(),
                    content: line.to_string(),
                });

                // If length > 100 messages, remove the oldest message (keep system prompt)
                if conversation_history.len() > 100 {
                    conversation_history.remove(1);
                }

                // Send message and get response
                match send_message(line.as_str(), &config, conversation_history.clone()).await {
                    Ok(response_content) => {
                        // Add assistant response to conversation history
                        conversation_history.push(Message {
                            role: "assistant".to_string(),
                            content: response_content.clone(),
                        });
                    }
                    Err(e) => {
                        eprintln!("{}", format!("Error: {}", e).red());
                        // Remove the failed message from history
                        if !conversation_history.is_empty() {
                            conversation_history.pop();
                        }
                        continue;
                    }
                }
            }
            Err(e) => {
                eprintln!("{}", format!("Error reading input: {}", e).red());
                break;
            }
        }
    }

    Ok(())
}

/// Handles special commands starting with /
/// Returns a tuple: (should_exit, new_config_option, message_to_send)
fn handle_command(
    command: &str,
    conversation_history: &mut Vec<Message>,
) -> Result<(bool, Option<ConfigFile>, Option<String>), Box<dyn Error>> {
    match command {
        "/exit" | "/quit" | "/q" => Ok((true, None, None)),
        "/help" | "/h" | "/?" | "/commands" | "/cmds" => {
            println!("{}", "Available commands:".bold());
            println!("  /help               - Show this help message");
            println!("  /exit, /quit, /q    - Exit the program");
            println!("  /profile <name>     - Switch to a different config profile");
            println!("  /clear              - Clear conversation history");
            println!("  /keys               - Open the API keys file");
            println!("  /listprofiles, /lsp - List all available profiles");
            println!("  /multiline, /ml, // - Enter multiline input mode");
            Ok((false, None, None))
        }
        "/clear" => {
            conversation_history.clear();
            println!("{}", "Conversation history cleared.".blue());
            Ok((false, None, None))
        }
        "/keys" | "/key" | "/apikeys" | "/apikey" => {
            open_keys()?;
            Ok((false, None, None))
        }
        cmd if cmd.starts_with("/profile") || cmd.starts_with("/p") => {
            let parts: Vec<&str> = command.split_whitespace().collect();
            if parts.len() != 2 {
                println!("Usage: /profile <profile_name>");
                return Ok((false, None, None));
            }

            let profile_name = parts[1];
            match load_config(Some(profile_name)) {
                Ok(new_config) => {
                    println!(
                        "{}",
                        format!("Switched to profile '{}'", profile_name).green()
                    );
                    Ok((false, Some(new_config), None))
                }
                Err(e) => {
                    eprintln!(
                        "{}",
                        format!("Failed to switch to profile '{}': {}", profile_name, e).red()
                    );
                    Ok((false, None, None))
                }
            }
        }
        "/listprofiles" | "/lsp" => {
            let path = get_config_path()?;
            let config = fs::read_to_string(&path)?;

            // Parse the JSON config to get root keys (profile names)
            let root_value: serde_json::Value = serde_json::from_str(&config)
                .map_err(|e| format!("Failed to parse config file: {}", e))?;

            println!("{}", "Available profiles:".bold());
            if let Some(obj) = root_value.as_object() {
                if obj.is_empty() {
                    println!("  No profiles found");
                } else {
                    for key in obj.keys() {
                        println!("  {}", key);
                    }
                }
            } else {
                eprintln!("  Invalid config format - expected JSON object");
            }
            Ok((false, None, None))
        }
        "/multiline" | "/ml" | "//" => {
            if let Some(multiline_input) = Editor::new().edit("").unwrap() {
                if !multiline_input.trim().is_empty() {
                    println!("{}\n", multiline_input);
                    // Return the multiline input as a message to be processed
                    return Ok((false, None, Some(multiline_input)));
                } else {
                    println!("{}", "Empty input - cancelled".yellow());
                }
            } else {
                eprintln!("{}", "Cancelled".red());
            }
            Ok((false, None, None))
        }
        _ => {
            eprintln!(
                "{}",
                format!(
                    "Unknown command: {}. Type /help for available commands.",
                    command
                )
                .red()
            );
            Ok((false, None, None))
        }
    }
}

async fn send_message(
    _message: &str,
    config: &Config,
    history: Vec<Message>,
) -> Result<String, Box<dyn Error>> {
    let spinner = ProgressBar::new_spinner();
    spinner.enable_steady_tick(std::time::Duration::from_millis(100));
    spinner.set_style(
        ProgressStyle::with_template("{spinner:.blue} {msg} {elapsed:.bold}")
            .unwrap()
            .tick_chars("⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏"),
    );
    spinner.set_message("Thinking...");

    // Build messages array with history
    let mut messages = vec![];

    if let Some(file) = &config.file {
        if let Some(system_prompt) = &file.system_prompt {
            messages.push(serde_json::json!({ "role": "system", "content": system_prompt }));
        }
    }

    for msg in history {
        messages.push(serde_json::json!({ "role": msg.role, "content": msg.content }));
    }

    let body = if let Some(file) = &config.file {
        serde_json::json!({
            "model": file.model,
            "messages": messages,
            "stream": config.stream,
        })
    } else {
        return Err("No configuration file found".to_string().into());
    };

    // Send the message
    let client = Client::new();

    let mut response;

    if config.verbose {
        dbg!(&body);
    }

    let (url, auth_header) = if let Some(file) = &config.file {
        match file.provider.as_str() {
            "pollinations" => ("https://text.pollinations.ai/openai", None),
            "google" => (
                "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
                Some(format!("Bearer {}", load_keys().unwrap().google.unwrap())),
            ),
            "openai" => (
                "https://api.openai.com/v1/chat/completions",
                Some(format!("Bearer {}", load_keys().unwrap().openai.unwrap())),
            ),
            provider => {
                return Err(format!("Unsupported provider: {}", provider).into());
            }
        }
    } else {
        return Err("No configuration file found".to_string().into());
    };

    let mut request = client
        .post(url)
        .header("Content-Type", "application/json")
        .body(serde_json::to_string(&body)?);

    if let Some(auth) = auth_header {
        request = request.header("Authorization", auth);
    }

    response = request.send().await?;

    if !response.status().is_success() {
        return Err(format!("Request failed with status: {}", response.status()).into());
    }

    spinner.finish_and_clear();

    if config.stream == Some(true) {
        spinner.finish_and_clear();
        let mut buffer = String::new();
        let mut response_content = String::new();
        let mut done = false;
        while let Some(chunk) = response.chunk().await? {
            let chunk_str = String::from_utf8_lossy(&chunk);
            buffer.push_str(&chunk_str);

            loop {
                if let Some(newline_idx) = buffer.find('\n') {
                    let line = buffer[..newline_idx].trim_end_matches('\r').to_string();
                    buffer.drain(..=newline_idx);
                    if line.is_empty() {
                        continue;
                    }

                    if line.starts_with("data: ") {
                        let payload = line[6..].trim();
                        if payload == "[DONE]" {
                            done = true;
                            println!(); // Fixes issue on Linux where % sign shows at end
                            break;
                        } else {
                            if let Ok(v) = serde_json::from_str::<serde_json::Value>(payload) {
                                if let Some(choices) = v.get("choices").and_then(|c| c.as_array()) {
                                    for choice in choices {
                                        if let Some(delta) = choice.get("delta") {
                                            if let Some(content) =
                                                delta.get("content").and_then(|c| c.as_str())
                                            {
                                                print!("{}", content);
                                                let _ = std::io::stdout().flush();
                                                response_content.push_str(content);
                                            }
                                        } else if let Some(content) = choice
                                            .get("message")
                                            .and_then(|m| m.get("content"))
                                            .and_then(|c| c.as_str())
                                        {
                                            print!("{}", content);
                                            let _ = std::io::stdout().flush();
                                            response_content.push_str(content);
                                        }
                                    }
                                }
                            }
                        }
                    } else if config.verbose {
                        eprintln!("{}", line);
                    }
                } else {
                    break;
                }
            }

            if done {
                break;
            }
        }
        Ok(response_content)
    } else {
        let response_text = response.text().await?;
        let json: serde_json::Value = serde_json::from_str(&response_text)?;
        let content = json["choices"][0]["message"]["content"]
            .as_str()
            .unwrap_or("")
            .replace("\\n", "\n")
            .trim_matches('"')
            .to_string();
        spinner.finish_and_clear();
        println!("\n{}\n", content);
        Ok(content)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_run() {
        let config = Config {
            verbose: false,
            interactive: false,
            message: Some(String::from("Hello, world!")),
            stream: Some(false),
            file: Some(ConfigFile {
                provider: String::from("pollinations"),
                model: String::from("openai"),
                stream: Some(false),
                system_prompt: Some(String::from("You are a helpful assistant named Tofu.")),
            }),
        };
        let result = run(config).await;
        assert!(result.is_ok());
    }
}