toast-api 0.1.3

use anyhow::{anyhow, Context, Result};
use ctrlc;
use regex::Regex;
use std::fs;
use std::io::{self, Write};
use std::path::Path;
use std::process::{self, Command, Stdio};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::future::Future;
use std::pin::Pin;

use crate::api::{Attachment, Claude, Session as ClaudeSession};
use crate::deepseek::{DeepSeek, Session as DeepSeekSession};
use crate::config::{HAIKU_MODEL, MAX_INTERNAL_ITERS, OPUS_MODEL, SONNET_MODEL, SYSTEM_PROMPT};
use crate::utils::{extract_commands, prettify};

/// Unified CLI arguments for both Claude and DeepSeek
#[derive(Debug)]
pub struct UnifiedArgs {
    pub use_deepseek: bool,
    pub use_opus: bool,
    pub use_haiku: bool,
}

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

/// Read files into attachments for Claude
fn collect_claude_attachments(paths: &[&str]) -> Result<Vec<Attachment>> {
    const LIMIT: usize = 5;
    const SIZE_LIMIT: u64 = 10 * 1024 * 1024;
    
    if paths.len() > LIMIT {
        return Err(anyhow!("cannot attach more than {LIMIT} files"));
    }
    
    let mut atts = Vec::new();
    
    for p in paths {
        if let Ok(meta) = fs::metadata(p) {
            if meta.len() > SIZE_LIMIT {
                eprintln!("Warning: file {p} is larger than 10 MB, skipping");
                continue;
            }
            
            if let Ok(content) = fs::read_to_string(p) {
                atts.push(Attachment {
                    file_name: Path::new(p)
                        .file_name()
                        .unwrap_or_default()
                        .to_string_lossy()
                        .into(),
                    size: meta.len(),
                    content,
                });
            } else {
                eprintln!("Warning: couldn't read file {p}");
            }
        } else {
            eprintln!("Warning: couldn't access file {p}");
        }
    }
    
    Ok(atts)
}

/// Run the unified CLI application with provider selection
pub async fn run(args: UnifiedArgs) -> Result<()> {
    // 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);
        })?;
    }
    
    if args.use_deepseek {
        run_deepseek(args, running).await
    } else {
        run_claude(args, running).await
    }
}

/// Run the CLI with DeepSeek provider
async fn run_deepseek(args: UnifiedArgs, running: Arc<AtomicBool>) -> Result<()> {
    // 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\nTo get your DeepSeek auth token:\n1. Go to chat.deepseek.com in your browser\n2. Open Developer Tools (F12)\n3. Go to Network tab\n4. Look for Authorization header in any request\n5. Save the token part (without 'Bearer ') to this file",
            auth_token_path
        ));
    };
    
    // 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\nDeepSeek requires Cloudflare cookies.\nUse the deepseek4free library to generate them.",
            cookies_path
        ));
    };
    
    let session = DeepSeekSession {
        auth_token,
        cookies,
    };
    
    // Determine model based on flags
    let model = if args.use_opus {
        "deepseek-coder"
    } else if args.use_haiku {
        "deepseek-lite"
    } else {
        "deepseek-chat" // Default model
    };
    
    let mut deepseek = DeepSeek::new_with_model(session, model)?;
    
    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 with DeepSeek!\n");
            id
        }
        Err(e) => {
            return Err(anyhow!("Failed to create DeepSeek chat session: {}", e));
        }
    };
    
    // Set thinking mode to disabled for now
    let thinking_mode = crate::deepseek::ThinkingMode::Disabled;
    let search_mode = crate::deepseek::SearchMode::Disabled;
    
    // 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;
        }
        
        // Handle exec commands
        if let Some(caps) = crate::utils::EXEC_RE.captures(input) {
            let cmd = caps[1].to_string();
            match execute_command(&cmd) {
                Ok(output) => {
                    let msg = format!("Command executed: {cmd}\n\n{output}");
                    print!("DeepSeek: ");
                    stdout.flush()?;
                    
                    match deepseek.chat_completion(&chat_id, &msg, None, thinking_mode.clone(), search_mode.clone()).await {
                        Ok(response) => {
                            println!("{}", prettify(&response));
                            process_deepseek_commands(&mut deepseek, &chat_id, &response, thinking_mode.clone(), search_mode.clone()).await?;
                        }
                        Err(e) => {
                            eprintln!("\nError: {}", e);
                        }
                    }
                }
                Err(e) => {
                    eprintln!("Warning: command execution failed: {e}");
                }
            }
            continue;
        }
        
        // Handle read_file commands
        if let Some(caps) = crate::utils::READ_RE.captures(input) {
            let paths: Vec<String> = caps[1].split_whitespace().map(String::from).collect();
            let mut file_contents = Vec::new();
            
            for path in &paths {
                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!("DeepSeek: ");
            stdout.flush()?;
            
            match deepseek.chat_completion(&chat_id, &file_message, None, thinking_mode.clone(), search_mode.clone()).await {
                Ok(response) => {
                    println!("{}", prettify(&response));
                    process_deepseek_commands(&mut deepseek, &chat_id, &response, thinking_mode.clone(), search_mode.clone()).await?;
                }
                Err(e) => {
                    eprintln!("\nError: {}", e);
                }
            }
            continue;
        }
        
        // Regular message
        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_deepseek_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_deepseek_commands(
    deepseek: &mut DeepSeek, 
    chat_id: &str, 
    response: &str,
    thinking_mode: crate::deepseek::ThinkingMode,
    search_mode: crate::deepseek::SearchMode,
) -> Result<()> {
    process_deepseek_commands_internal(deepseek, chat_id, response, thinking_mode, search_mode, 0).await
}

fn process_deepseek_commands_internal<'a>(
    deepseek: &'a mut DeepSeek, 
    chat_id: &'a str, 
    response: &'a str,
    thinking_mode: crate::deepseek::ThinkingMode,
    search_mode: crate::deepseek::SearchMode,
    depth: usize,
) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
    Box::pin(async move {
        // Avoid infinite recursion
        if depth >= MAX_INTERNAL_ITERS {
            println!("Max internal iterations reached, returning to user.");
            return Ok(());
        }
        
        let (reads, execs) = extract_commands(response);
        if reads.is_empty() && execs.is_empty() {
            return Ok(());
        }
        
        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"));
            
            match deepseek.chat_completion(chat_id, &file_message, None, thinking_mode.clone(), search_mode.clone()).await {
                Ok(resp) => {
                    println!("DeepSeek: {}", prettify(&resp));
                    return process_deepseek_commands_internal(deepseek, chat_id, &resp, thinking_mode, search_mode, depth + 1).await;
                }
                Err(e) => {
                    eprintln!("Error: {}", e);
                    return Err(e);
                }
            }
        }
        
        if !execs.is_empty() {
            let mut outputs = String::new();
            
            for cmd in &execs {
                match execute_command(cmd) {
                    Ok(output) => outputs.push_str(&output),
                    Err(e) => outputs.push_str(&format!("Command execution failed: {e}")),
                }
                outputs.push_str("\n\n---\n\n");
            }
            
            match deepseek.chat_completion(chat_id, &outputs, None, thinking_mode.clone(), search_mode.clone()).await {
                Ok(resp) => {
                    println!("DeepSeek: {}", prettify(&resp));
                    return process_deepseek_commands_internal(deepseek, chat_id, &resp, thinking_mode, search_mode, depth + 1).await;
                }
                Err(e) => {
                    eprintln!("Error: {}", e);
                    return Err(e);
                }
            }
        }
        
        Ok(())
    })
}

/// Extract organization ID from cookie string
fn extract_org_id_from_cookie(cookie: &str) -> Option<String> {
    let re = Regex::new(r"lastActiveOrg=([0-9a-f-]+)").ok()?;
    re.captures(cookie)
        .and_then(|caps| caps.get(1))
        .map(|m| m.as_str().to_string())
}

/// Run the CLI with Claude provider
async fn run_claude(args: UnifiedArgs, running: Arc<AtomicBool>) -> Result<()> {
    // Load session values from config files
    let config_dir = dirs::config_dir()
        .ok_or_else(|| anyhow!("Could not determine config directory"))?
        .join("toast");

    let cookie_path = config_dir.join("cookie");
    let org_id_path = config_dir.join("org_id");

    // Check if config directory exists, if not create it and provide instructions
    if !config_dir.exists() {
        fs::create_dir_all(&config_dir).context(format!(
            "Failed to create config directory at {:?}",
            config_dir
        ))?;
        return Err(anyhow!(
            "Configuration directory created at {:?}\n\nPlease create a cookie file with your Claude cookie", 
            config_dir,
        ));
    }

    // Check and load cookie
    let cookie = if cookie_path.exists() {
        fs::read_to_string(&cookie_path)
            .context(format!("Failed to read cookie from {:?}", cookie_path))?
            .trim()
            .to_string()
    } else {
        return Err(anyhow!(
            "Cookie file not found at {:?}",
            cookie_path,
        ));
    };

    // Check and load org_id, or extract from cookie if file doesn't exist
    let org_id = if org_id_path.exists() {
        fs::read_to_string(&org_id_path)
            .context(format!(
                "Failed to read organization ID from {:?}",
                org_id_path
            ))?
            .trim()
            .to_string()
    } else {
        // Try to extract org_id from cookie
        if let Some(extracted_org_id) = extract_org_id_from_cookie(&cookie) {
            // Save the extracted org_id to the file for future use
            fs::write(&org_id_path, &extracted_org_id).context(format!(
                "Failed to write organization ID to {:?}",
                org_id_path
            ))?;
            println!(
                "Extracted organization ID from cookie and saved to {:?}",
                org_id_path
            );
            extracted_org_id
        } else {
            return Err(anyhow!(
                "Organization ID file not found at {:?} and couldn't extract it from cookie.",
                org_id_path,
            ));
        }
    };

    let user_agent =
        "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:137.0) Gecko/20100101 Firefox/137.0"
            .to_string();

    let session = ClaudeSession {
        cookie,
        user_agent,
        organization_id: org_id,
    };

    // Determine model based on flags
    let model: &str = if args.use_opus {
        OPUS_MODEL
    } else if args.use_haiku {
        HAIKU_MODEL
    } else {
        SONNET_MODEL
    };

    let claude = Claude::new(session.clone(), model)?;
    println!("Starting new Claude chat session using model: {}", model);

    let stdin = io::stdin();
    let mut stdout = io::stdout();
    let mut chat_id = String::new();
    let mut system_prompt_sent = false;

    while running.load(Ordering::SeqCst) {
        print!("You: ");
        stdout.flush()?;
        let mut buf = String::new();
        stdin.read_line(&mut buf)?;
        let input = buf.trim_end();
        if input == "" {
            continue;
        }
        if input.eq_ignore_ascii_case("/exit") || input.eq_ignore_ascii_case("exit") || input == "x"
        {
            if !chat_id.is_empty() {
                claude.delete_chat(&chat_id).await.ok();
            }
            break;
        }

        // Initialize chat
        if chat_id.is_empty() {
            chat_id = claude.create_chat().await.context("creating chat")?;
        }

        // Handle exec commands
        if let Some(caps) = crate::utils::EXEC_RE.captures(input) {
            let cmd = caps[1].to_string();
            if !system_prompt_sent {
                claude
                    .send_message(&chat_id, SYSTEM_PROMPT, &[])
                    .await
                    .context("sending system prompt")?;
                system_prompt_sent = true;
            }
            
            match execute_command(&cmd) {
                Ok(output) => {
                    let msg = format!("Command executed: {cmd}\n\n{output}");
                    let ans = claude.send_message(&chat_id, &msg, &[]).await?;
                    println!("Claude:\n{}", prettify(&ans));
                    process_claude_commands(&claude, &chat_id, &ans).await?;
                }
                Err(e) => {
                    eprintln!("Warning: command execution failed: {e}");
                    let msg = format!("Command execution failed: {e}");
                    let ans = claude.send_message(&chat_id, &msg, &[]).await?;
                    println!("Claude:\n{}", prettify(&ans));
                }
            }
            continue;
        }

        // Handle read_file commands
        if let Some(caps) = crate::utils::READ_RE.captures(input) {
            let paths: Vec<String> = caps[1].split_whitespace().map(String::from).collect();
            let path_refs: Vec<&str> = paths.iter().map(String::as_str).collect();
            if !system_prompt_sent {
                claude
                    .send_message(&chat_id, SYSTEM_PROMPT, &[])
                    .await
                    .context("sending system prompt")?;
                system_prompt_sent = true;
            }
            
            let rest = input.strip_prefix(&caps[0]).unwrap_or("").trim();
            let attachments = collect_claude_attachments(&path_refs).unwrap_or_default();
            let ans = claude
                .send_message(&chat_id, rest, &attachments)
                .await
                .context("sending user message")?;
                
            println!("Claude:\n{}", prettify(&ans));
            process_claude_commands(&claude, &chat_id, &ans).await?;
        } else {
            // Regular message
            if !system_prompt_sent {
                claude
                    .send_message(&chat_id, SYSTEM_PROMPT, &[])
                    .await
                    .context("sending system prompt")?;
                system_prompt_sent = true;
            }
            
            let ans = claude
                .send_message(&chat_id, input, &[])
                .await
                .context("sending user message")?;
                
            println!("Claude:\n{}", prettify(&ans));
            process_claude_commands(&claude, &chat_id, &ans).await?;
        }
    }
    
    Ok(())
}

/// Process Claude's responses for internal tool commands
async fn process_claude_commands(claude: &Claude, chat_id: &str, response: &str) -> Result<()> {
    process_claude_commands_internal(claude, chat_id, response, 0).await
}

fn process_claude_commands_internal<'a>(claude: &'a Claude, chat_id: &'a str, response: &'a str, depth: usize) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
    Box::pin(async move {
        // Avoid infinite recursion
        if depth >= MAX_INTERNAL_ITERS {
            println!("Max internal iterations reached, returning to user.");
            return Ok(());
        }
        
        let (reads, execs) = extract_commands(response);
        if reads.is_empty() && execs.is_empty() {
            return Ok(());
        }
        
        if !reads.is_empty() {
            let atts = collect_claude_attachments(&reads.iter().map(String::as_str).collect::<Vec<_>>())
                .unwrap_or_default();
                
            match claude.send_message(chat_id, "read_file response:", &atts).await {
                Ok(resp) => {
                    println!("Claude:\n{}", prettify(&resp));
                    return process_claude_commands_internal(claude, chat_id, &resp, depth + 1).await;
                }
                Err(e) => {
                    return Err(e);
                }
            }
        }
        
        if !execs.is_empty() {
            let mut outputs = String::new();
            
            for cmd in &execs {
                match execute_command(cmd) {
                    Ok(output) => outputs.push_str(&output),
                    Err(e) => outputs.push_str(&format!("Command execution failed: {e}")),
                }
                outputs.push_str("\n\n---\n\n");
            }
            
            match claude.send_message(chat_id, &outputs, &[]).await {
                Ok(resp) => {
                    println!("Claude:\n{}", prettify(&resp));
                    return process_claude_commands_internal(claude, chat_id, &resp, depth + 1).await;
                }
                Err(e) => {
                    return Err(e);
                }
            }
        }
        
        Ok(())
    })
}