toast-api 0.1.4

use anyhow::{anyhow, Context, Result};
use clap::Parser;
use ctrlc;
use std::fs;
use std::io::{self, stdout, Write};
use std::process;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::time::sleep;
use std::future::Future;
use std::pin::Pin;

use crate::deepseek::{DeepSeek, Session, SearchMode, ThinkingMode};
// No need to import from config as we're using string literals directly
use crate::utils::{extract_commands, prettify};

/// CLI arguments for the DeepSeek interface
#[derive(Parser, Debug)]
#[clap(name = "deepseek", about = "DeepSeek CLI client")]
struct Args {
    /// Enable web search
    #[clap(long)]
    search: bool,
    
    /// Disable thinking mode
    #[clap(long)]
    no_thinking: bool,

    /// Use opus model
    #[clap(long, conflicts_with = "haiku")]
    opus: bool,

    /// Use haiku model
    #[clap(long, conflicts_with = "opus")]
    haiku: bool,
}

/// Run the DeepSeek CLI application
pub async fn run() -> Result<()> {
    let args = Args::parse();
    
    // Load session values from config files
    let config_dir = dirs::config_dir()
        .ok_or_else(|| anyhow!("Could not determine config directory"))?
        .join("toast")
        .join("deepseek");
    
    // Create config directory if it doesn't exist
    if !config_dir.exists() {
        fs::create_dir_all(&config_dir)?;
    }
    
    let auth_token_path = config_dir.join("auth_token");
    let cookies_path = config_dir.join("cookies.json");
    
    // Check auth token
    let auth_token = if auth_token_path.exists() {
        fs::read_to_string(&auth_token_path)
            .context(format!("Failed to read auth token from {:?}", auth_token_path))?
            .trim()
            .to_string()
    } else {
        return Err(anyhow!(
            "Auth token file not found at {:?}\n\n{}",
            auth_token_path,
            get_config_help("deepseek_auth_token")
        ));
    };
    
    // Check cookies
    let cookies = if cookies_path.exists() {
        serde_json::from_str(&fs::read_to_string(&cookies_path)
            .context(format!("Failed to read cookies from {:?}", cookies_path))?)?
    } else {
        return Err(anyhow!(
            "Cookies file not found at {:?}\n\n{}",
            cookies_path,
            get_config_help("deepseek_cookies")
        ));
    };
    
    let session = Session {
        auth_token,
        cookies,
    };
    
    let thinking_mode = if args.no_thinking {
        ThinkingMode::Disabled
    } else {
        ThinkingMode::Detailed
    };
    
    let search_mode = if args.search {
        SearchMode::Enabled
    } else {
        SearchMode::Disabled
    };
    
    // Determine model based on flags
    let model = if args.opus {
        "deepseek-coder"
    } else if args.haiku {
        "deepseek-lite"
    } else {
        "deepseek-chat" // Default model
    };
    
    let mut deepseek = DeepSeek::new_with_model(session, model)?;
    
    // Set up Ctrl-C handler
    let running = Arc::new(AtomicBool::new(true));
    {
        let running = running.clone();
        ctrlc::set_handler(move || {
            running.store(false, Ordering::SeqCst);
            println!("\nGoodbye!");
            process::exit(0);
        })?;
    }
    
    let stdin = io::stdin();
    let mut stdout = io::stdout();
    
    // Create a new chat session
    println!("Starting new DeepSeek chat session...");
    let chat_id = match deepseek.create_chat_session().await {
        Ok(id) => {
            println!("Session started!\n");
            id
        }
        Err(e) => {
            return Err(anyhow!("Failed to create chat session: {}", e));
        }
    };
    
    // Main chat loop
    while running.load(Ordering::SeqCst) {
        print!("You: ");
        stdout.flush()?;
        
        let mut buf = String::new();
        stdin.read_line(&mut buf)?;
        let input = buf.trim_end();
        
        // Check for empty input or exit commands
        if input.is_empty() {
            continue;
        }
        
        if input.eq_ignore_ascii_case("/exit") || input.eq_ignore_ascii_case("exit") || input == "x" {
            break;
        }
        
        // Send message to API
        print!("DeepSeek: ");
        stdout.flush()?;
        
        match deepseek.chat_completion(&chat_id, input, None, thinking_mode.clone(), search_mode.clone()).await {
            Ok(response) => {
                println!("{}", prettify(&response));
                
                // Process commands in the response
                process_commands(&mut deepseek, &chat_id, &response, thinking_mode.clone(), search_mode.clone()).await?;
            }
            Err(e) => {
                eprintln!("\nError: {}", e);
            }
        }
        println!();
    }
    
    Ok(())
}

// Process commands in DeepSeek's response
async fn process_commands(
    deepseek: &mut DeepSeek,
    chat_id: &str,
    response: &str,
    thinking_mode: ThinkingMode,
    search_mode: SearchMode,
) -> Result<()> {
    // Handle commands with max depth to avoid infinite recursion
    process_commands_internal(deepseek, chat_id, response, thinking_mode, search_mode, 0).await
}

// Internal implementation with max depth control using Pin<Box<dyn Future>>
fn process_commands_internal<'a>(
    deepseek: &'a mut DeepSeek,
    chat_id: &'a str,
    response: &'a str,
    thinking_mode: ThinkingMode,
    search_mode: SearchMode,
    depth: u8,
) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
    Box::pin(async move {
    // Limit recursion depth
    const MAX_DEPTH: u8 = 5;
    if depth >= MAX_DEPTH {
        println!("Maximum command processing depth reached ({}). Stopping recursion.", MAX_DEPTH);
        return Ok(());
    }
    
    // Extract read_file and exec commands
    let (reads, execs) = extract_commands(response);
    
    if reads.is_empty() && execs.is_empty() {
        return Ok(());
    }
    
    // Short pause before processing commands
    sleep(Duration::from_millis(500)).await;
    
    // Process file reads
    if !reads.is_empty() {
        let mut file_contents = Vec::new();
        
        for path in &reads {
            match fs::read_to_string(path) {
                Ok(content) => {
                    file_contents.push(format!("=== File: {} ===\n{}", path, content));
                }
                Err(e) => {
                    file_contents.push(format!("Error reading file {}: {}", path, e));
                }
            }
        }
        
        let file_message = format!("Here are the contents of the files you requested:\n\n{}", 
                                  file_contents.join("\n\n"));
        
        print!("Sending file contents... ");
        stdout().flush()?;
        
        match deepseek.chat_completion(chat_id, &file_message, None, thinking_mode, search_mode).await {
            Ok(response) => {
                println!("Done!");
                println!("DeepSeek: {}", prettify(&response));
                
                // Process next level of commands
                process_commands_internal(deepseek, chat_id, &response, thinking_mode, search_mode, depth + 1).await?;
            }
            Err(e) => {
                println!("Error: {}", e);
            }
        }
    }
    
    // Process exec commands
    if !execs.is_empty() {
        for cmd in &execs {
            println!("\nExecuting: {}", cmd);
            
            match execute_command(cmd) {
                Ok(output) => {
                    println!("{}", output);
                    
                    print!("Sending command results... ");
                    stdout().flush()?;
                    
                    let cmd_message = format!("Command executed: {}\n\nOutput:\n{}", cmd, output);
                    
                    match deepseek.chat_completion(chat_id, &cmd_message, None, thinking_mode, search_mode).await {
                        Ok(response) => {
                            println!("Done!");
                            println!("DeepSeek: {}", prettify(&response));
                            
                            // Process next level of commands
                            process_commands_internal(deepseek, chat_id, &response, thinking_mode, search_mode, depth + 1).await?;
                        }
                        Err(e) => {
                            println!("Error: {}", e);
                        }
                    }
                }
                Err(e) => {
                    println!("Error executing command: {}", e);
                }
            }
        }
    }
    
    Ok(())
    }) // Close the Box::pin(async move {}) block
}

// Execute a shell command and return the output
fn execute_command(command: &str) -> Result<String> {
    let output = process::Command::new("sh")
        .arg("-c")
        .arg(command)
        .output()?;
    
    let mut result = String::new();
    
    if !output.stdout.is_empty() {
        result.push_str("=== STDOUT ===\n");
        result.push_str(&String::from_utf8_lossy(&output.stdout));
        result.push('\n');
    }
    
    if !output.stderr.is_empty() {
        result.push_str("=== STDERR ===\n");
        result.push_str(&String::from_utf8_lossy(&output.stderr));
        result.push('\n');
    }
    
    result.push_str(&format!("Exit code: {}", output.status.code().unwrap_or(-1)));
    
    Ok(result)
}

// Configuration help text
fn get_config_help(file_name: &str) -> String {
    match file_name {
        "deepseek_auth_token" => "To get your DeepSeek auth token:
1. Go to chat.deepseek.com in your browser
2. Log in to your account
3. Open Developer Tools (F12 or right-click and select 'Inspect')
4. Go to the Network tab
5. Refresh the page or make a request
6. Look for requests to the DeepSeek API
7. In the 'Headers' tab, find 'Request Headers'
8. Look for the 'Authorization' header with format 'Bearer {token}'
9. Copy the token part (without 'Bearer ') and save it to this folder with filename: auth_token".to_string(),
        
        "deepseek_cookies" => "The DeepSeek API requires Cloudflare cookies to bypass protection:
1. Run the Python script from the deepseek4free/dsk folder:
   python bypass.py
2. Copy the generated cookies.json file to this folder".to_string(),
        
        _ => format!("Configuration file {} is missing.", file_name),
    }
}