toast-api 0.1.9

use anyhow::{anyhow, Context, Result};
use clap::Parser;
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 crate::api::{Attachment, Claude, Session};
use crate::config::{HAIKU_MODEL, MAX_INTERNAL_ITERS, OPUS_MODEL, SONNET_MODEL, SYSTEM_PROMPT};
use crate::utils::{extract_commands, prettify};

/// CLI arguments for the default interface
#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
    /// Use opus model
    #[clap(long, conflicts_with_all = ["haiku", "custom_model"])]
    opus: bool,

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

    /// Specify a custom model ID
    #[clap(long, conflicts_with_all = ["opus", "haiku"])]
    custom_model: Option<String>,
}

/// Run the CLI application
pub fn run() -> Result<()> {
    tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .unwrap()
        .block_on(async_main())
}

fn get_config_help(file_name: &str) -> String {
    let cookie_help = "To get your cookie:
1. Go to claude.ai in your browser
2. Open Developer Tools (F12 or right-click and select 'Inspect')
3. Go to the Network tab
4. Refresh the page
5. Click on any request to claude.ai
6. In the 'Headers' tab, find 'Request Headers'
7. Look for the 'Cookie' header
8. Copy the entire cookie value and save it to this folder with filename: cookie";

    match file_name {
        "cookie" => cookie_help.to_string(),
        _ => format!("Configuration file {} is missing.", file_name),
    }
}

/// Extract organization ID from cookie string
pub 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())
}

pub async fn async_main() -> 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");

    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 the following files:\n\n1. cookie file:\n{}\n\n2", 
            config_dir,
            get_config_help("cookie"),
        ));
    }

    // 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 {:?}\n\n{}",
            cookie_path,
            get_config_help("cookie")
        ));
    };

    // 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.\n\n{}",
                org_id_path,
                get_config_help("org_id")
            ));
        }
    };

    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 = Session {
        cookie,
        user_agent,
        organization_id: org_id,
    };

    // Determine model based on flags
    let model: &str = if let Some(custom) = args.custom_model {
        Box::leak(custom.into_boxed_str())
    } else if args.opus {
        OPUS_MODEL
    } else if args.haiku {
        HAIKU_MODEL
    } else {
        SONNET_MODEL
    };
    let claude = Claude::new(session.clone(), model)?;

    // 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();
    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;
            }
            run_exec(&claude, &chat_id, &cmd).await?;
            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_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(&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(&claude, &chat_id, ans).await?;
        }
    }
    Ok(())
}

/// Execute shell command and send output to Claude
async fn run_exec(claude: &Claude, chat_id: &str, cmd: &str) -> Result<()> {
    let out = match execute_command(cmd) {
        Ok(output) => output,
        Err(e) => {
            eprintln!("Warning: command execution failed: {e}");
            format!("Command execution failed: {e}")
        }
    };
    let msg = format!("Command executed: {cmd}\n\n{out}");
    let ans = claude.send_message(chat_id, &msg, &[]).await?;
    println!("Claude:\n{}", prettify(&ans));
    process_claude(claude, chat_id, ans).await
}

/// 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
fn collect_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)
}

/// Process Claude's responses for internal tool commands
async fn process_claude(claude: &Claude, chat_id: &str, mut ans: String) -> Result<()> {
    for _ in 0..MAX_INTERNAL_ITERS {
        let (reads, execs) = extract_commands(&ans);
        if reads.is_empty() && execs.is_empty() {
            return Ok(());
        }
        if !reads.is_empty() {
            let atts = collect_attachments(&reads.iter().map(String::as_str).collect::<Vec<_>>())
                .unwrap_or_default();
            let ans2 = claude
                .send_message(chat_id, "read_file response:", &atts)
                .await?;
            println!("Claude:\n{}", prettify(&ans2));
            ans = ans2;
            continue;
        }
        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");
            }
            let ans2 = claude.send_message(chat_id, &outputs, &[]).await?;
            println!("Claude:\n{}", prettify(&ans2));
            ans = ans2;
            continue;
        }
    }
    println!("Max internal iterations reached, returning to user.");
    Ok(())
}