brokecrew/src/main.rs

mod cli;

use brokecrew::budget::BudgetTracker;
use brokecrew::client::OpenRouterClient;
use brokecrew::config::AppConfig;
use brokecrew::persona::Tier;
use brokecrew::pipeline::{PipelineConfig, PipelineEvent, PipelineExecutor};
use brokecrew::registry::ModelRegistry;
use brokecrew::router::{RouteStage, Router};
use clap::Parser;
use cli::{ChatArgs, Cli, Commands, ConfigSubcommand, PersonasSubcommand};
use indicatif::{ProgressBar, ProgressStyle};
use owo_colors::OwoColorize;
use std::io::{self, BufRead, Write};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::mpsc;
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};

fn init_logging(json_logs: bool, verbose: bool) {
    let filter = if verbose {
        EnvFilter::new("brokecrew=debug,reqwest=debug")
    } else {
        EnvFilter::new("brokecrew=info")
    };

    if json_logs {
        tracing_subscriber::registry()
            .with(filter)
            .with(tracing_subscriber::fmt::layer().json())
            .init();
    } else {
        tracing_subscriber::registry()
            .with(filter)
            .with(
                tracing_subscriber::fmt::layer()
                    .with_writer(io::stderr)
                    .without_time()
                    .with_target(false),
            )
            .init();
    }
}

fn get_persona_color_tag(persona: &str) -> String {
    match persona.to_lowercase().as_str() {
        "coder" => format!("[{}]", persona).bright_green().bold().to_string(),
        "reviewer" => format!("[{}]", persona).bright_yellow().bold().to_string(),
        "architect" => format!("[{}]", persona).bright_blue().bold().to_string(),
        "general" | "assistant" => format!("[{}]", persona).bright_magenta().bold().to_string(),
        _ => format!("[{}]", persona).bright_cyan().bold().to_string(),
    }
}

fn create_spinner(message: &'static str) -> ProgressBar {
    let pb = ProgressBar::new_spinner();
    pb.set_style(
        ProgressStyle::default_spinner()
            .tick_chars("⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏ ")
            .template("{spinner:.cyan} {msg}")
            .expect("valid template"),
    );
    pb.set_message(message);
    pb.enable_steady_tick(Duration::from_millis(80));
    pb
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let cli = Cli::parse();
    init_logging(cli.json_logs, cli.verbose);

    let config = AppConfig::load_from_paths(
        cli.config_file.as_deref(),
        cli.personas_file.as_deref(),
        cli.api_key.clone(),
    )
    .map_err(|e| format!("Configuration error: {e}"))?;

    let budget_path = config.config_dir.join("budget.json");
    let mut budget_tracker = BudgetTracker::new(Some(budget_path), config.credit_tier);
    let model_registry = Arc::new(ModelRegistry::new());

    // Check subcommand
    match cli.command {
        Some(Commands::Config(config_args)) => match config_args.action {
            ConfigSubcommand::SetKey { key } => {
                let api_key = match key {
                    Some(k) if !k.trim().is_empty() => k.trim().to_string(),
                    _ => {
                        eprint!("Enter OpenRouter API Key: ");
                        io::stderr().flush()?;
                        let mut input = String::new();
                        io::stdin().read_line(&mut input)?;
                        let trimmed = input.trim().to_string();
                        if trimmed.is_empty() {
                            eprintln!("{}", "Error: API key cannot be empty.".red().bold());
                            std::process::exit(1);
                        }
                        trimmed
                    }
                };

                AppConfig::save_api_key_to_keyring(&api_key)
                    .map_err(|e| format!("Failed to save API key to OS credential store: {e}"))?;
                println!(
                    "{}",
                    "✔ OpenRouter API key saved securely in OS credential store."
                        .bright_green()
                        .bold()
                );
                return Ok(());
            }
        },

        Some(Commands::Personas(personas_args)) => match personas_args.action {
            PersonasSubcommand::List => {
                println!("\n{}", "BrokeCrew Personas:".bold().underline());
                for persona in config.persona_registry.list() {
                    let badge = get_persona_color_tag(&persona.name);
                    println!("\n  {} - {}", badge, persona.system_prompt.dimmed());
                    if !persona.trigger_words.is_empty() {
                        println!(
                            "    Triggers: {}",
                            persona.trigger_words.join(", ").italic()
                        );
                    }
                    println!("    Candidates:");
                    for (i, c) in persona.candidates.iter().enumerate() {
                        let tier_tag = match c.tier {
                            Tier::Free => "FREE".bright_green().to_string(),
                            Tier::Paid => "PAID".bright_yellow().to_string(),
                        };
                        let stealth_tag = if c.is_stealth {
                            format!(" {}", "[stealth]".bright_magenta())
                        } else {
                            String::new()
                        };
                        println!(
                            "      {}. {} ({}){}",
                            i + 1,
                            c.id.bold(),
                            tier_tag,
                            stealth_tag
                        );
                    }
                }

                if !config.pipelines.is_empty() {
                    println!("\n{}", "Configured Pipelines:".bold().underline());
                    for pipe in &config.pipelines {
                        let steps_str: Vec<String> = pipe
                            .steps
                            .iter()
                            .map(|s| {
                                if s.pass_previous_output {
                                    format!("{} (with context)", s.persona.bold())
                                } else {
                                    s.persona.bold().to_string()
                                }
                            })
                            .collect();
                        println!(
                            "  • {} ➔ {}",
                            pipe.name.bright_cyan().bold(),
                            steps_str.join(" ➔ ")
                        );
                    }
                }
                println!();
                return Ok(());
            }
        },

        Some(Commands::Status) => {
            println!("\n{}", "BrokeCrew Status & Rate Limits:".bold().underline());
            // Fetch live registry for status count
            let spinner = create_spinner("Fetching live OpenRouter model catalog...");
            let model_count = match model_registry.fetch_live(&config.openrouter_base_url).await {
                Ok(count) => {
                    spinner.finish_and_clear();
                    count
                }
                Err(e) => {
                    spinner.finish_and_clear();
                    println!("  Live catalog: {}", format!("Offline / Error: {e}").red());
                    0
                }
            };

            if model_count > 0 {
                println!(
                    "  Live catalog: {} available models cached",
                    model_count.to_string().bright_green().bold()
                );
            }

            match budget_tracker.check_budget() {
                Ok(status) => {
                    let pct = (status.daily_requests * 100) / status.daily_limit;
                    let daily_color = if pct >= 80 {
                        format!(
                            "{}/{} ({}%)",
                            status.daily_requests, status.daily_limit, pct
                        )
                        .yellow()
                        .bold()
                        .to_string()
                    } else {
                        format!(
                            "{}/{} ({}%)",
                            status.daily_requests, status.daily_limit, pct
                        )
                        .green()
                        .bold()
                        .to_string()
                    };

                    println!("  Account tier: {:?}", config.credit_tier);
                    println!(
                        "  Rolling 1-min requests: {}/{}",
                        status.minute_requests.to_string().bold(),
                        status.minute_limit
                    );
                    println!("  Daily requests today: {}", daily_color);
                    if status.approaching_daily_limit {
                        println!("  {}", "⚠ Notice: Approaching daily free-tier limit. Non-urgent calls may be deferred.".yellow().bold());
                    }
                }
                Err(e) => {
                    println!("  Budget status: {}", e.to_string().red().bold());
                }
            }
            println!();
            return Ok(());
        }

        Some(Commands::Chat(args)) => {
            run_chat(args, &config, &mut budget_tracker, model_registry).await?;
        }

        None => {
            // Default to chat if no subcommand passed
            run_chat(
                ChatArgs::default(),
                &config,
                &mut budget_tracker,
                model_registry,
            )
            .await?;
        }
    }

    Ok(())
}

async fn run_chat(
    args: ChatArgs,
    config: &AppConfig,
    budget_tracker: &mut BudgetTracker,
    model_registry: Arc<ModelRegistry>,
) -> Result<(), Box<dyn std::error::Error>> {
    let api_key = match config.api_key.clone() {
        Some(k) if !k.trim().is_empty() => k,
        _ => {
            eprintln!(
                "{}\nRun: {}\nOr set: {}",
                "Error: OpenRouter API key not configured.".red().bold(),
                "brokecrew config set-key".bright_green().bold(),
                "export OPENROUTER_API_KEY=\"sk-or-v1-...\"".cyan()
            );
            std::process::exit(1);
        }
    };

    let client = OpenRouterClient::with_base_url(api_key, config.openrouter_base_url.clone());

    // Start background sync for model registry
    let reg_clone = model_registry.clone();
    let base_url = config.openrouter_base_url.clone();
    tokio::spawn(async move {
        let _ = reg_clone.fetch_live(&base_url).await;
    });

    let router = Router::new(
        config.default_persona.clone(),
        config.router_model.clone(),
        config.complexity_threshold,
    );

    // If prompt passed directly as CLI argument, execute single-shot
    if let Some(ref prompt) = args.prompt {
        execute_prompt(
            prompt,
            &args,
            config,
            budget_tracker,
            &model_registry,
            &router,
            &client,
        )
        .await?;
        return Ok(());
    }

    // Otherwise, run interactive REPL loop
    println!(
        "\n{}",
        "BrokeCrew Interactive Team Chat".bright_cyan().bold()
    );
    println!(
        "{}",
        "Type your request and press Enter. Type 'exit' or Ctrl+C to quit.\n".dimmed()
    );

    let stdin = io::stdin();
    let mut reader = stdin.lock();

    loop {
        print!("{} ", "you ❯".bright_white().bold());
        io::stdout().flush()?;

        let mut line = String::new();
        if reader.read_line(&mut line)? == 0 {
            break;
        }

        let trimmed = line.trim();
        if trimmed.is_empty() {
            continue;
        }
        if trimmed.eq_ignore_ascii_case("exit") || trimmed.eq_ignore_ascii_case("quit") {
            break;
        }

        execute_prompt(
            trimmed,
            &args,
            config,
            budget_tracker,
            &model_registry,
            &router,
            &client,
        )
        .await?;
        println!();
    }

    Ok(())
}

async fn execute_prompt(
    prompt: &str,
    args: &ChatArgs,
    config: &AppConfig,
    budget_tracker: &mut BudgetTracker,
    model_registry: &ModelRegistry,
    router: &Router,
    client: &OpenRouterClient,
) -> Result<(), Box<dyn std::error::Error>> {
    // 1. Budget check
    match budget_tracker.check_budget() {
        Ok(status) => {
            if status.approaching_daily_limit {
                eprintln!(
                    "{} Daily usage: {}/{}",
                    "⚠ Budget notice: Approaching daily ceiling.".yellow(),
                    status.daily_requests,
                    status.daily_limit
                );
            }
        }
        Err(e) => {
            eprintln!("{} {}", "Rate limit warning:".red().bold(), e);
            return Ok(());
        }
    }

    let allow_stealth = args.allow_stealth || config.allow_stealth;
    let require_paid = args.escalate;

    // 2. Check if running a multi-agent pipeline
    if let Some(ref pipe_name) = args.pipeline {
        let pipeline = config
            .pipelines
            .iter()
            .find(|p| p.name == *pipe_name)
            .ok_or_else(|| format!("Pipeline '{pipe_name}' not found in configuration"))?;

        run_pipeline(
            pipeline,
            config,
            budget_tracker,
            model_registry,
            client,
            prompt,
            allow_stealth,
            require_paid,
        )
        .await?;
        return Ok(());
    }

    // 3. Single persona dispatch
    let (target_persona, route_info) = if let Some(ref p) = args.persona {
        (p.clone(), None)
    } else {
        let spinner = create_spinner("Routing message...");
        // Record budget call if LLM router is invoked
        let decision = router
            .route(&config.persona_registry, client, prompt, require_paid)
            .await;
        spinner.finish_and_clear();

        match decision {
            Ok(dec) => {
                if dec.stage_used == RouteStage::StageBLlm {
                    let _ = budget_tracker.record_call(false);
                }
                let info = match dec.stage_used {
                    RouteStage::StageAKeyword => "routed via keyword pre-filter".to_string(),
                    RouteStage::StageBLlm => "routed via LLM classifier".to_string(),
                };
                (dec.persona, Some(info))
            }
            Err(e) => {
                eprintln!(
                    "Router warning: {}, falling back to {}",
                    e, config.default_persona
                );
                (config.default_persona.clone(), None)
            }
        }
    };

    if let Some(info) = route_info {
        eprintln!(
            "{}",
            format!("⚡ [BrokeCrew: {info} ➔ {}]", target_persona).dimmed()
        );
    }

    let (chunk_tx, mut chunk_rx) = mpsc::channel::<String>(100);
    let persona_tag = get_persona_color_tag(&target_persona);
    let spinner = create_spinner("Thinking...");

    let dispatch_future = PipelineExecutor::dispatch_persona(
        &target_persona,
        &config.persona_registry,
        model_registry,
        client,
        prompt,
        allow_stealth,
        require_paid,
        Some(chunk_tx),
    );

    let printer = tokio::spawn(async move {
        let mut first = true;
        while let Some(chunk) = chunk_rx.recv().await {
            if first {
                spinner.finish_and_clear();
                println!("\n{}", persona_tag);
                first = false;
            }
            print!("{chunk}");
            let _ = io::stdout().flush();
        }
    });

    let dispatch_result = dispatch_future.await;
    let _ = printer.await;

    match dispatch_result {
        Ok(res) => {
            println!();
            if res.fell_back {
                eprintln!(
                    "{}",
                    format!("⚠ Candidate fallback: served by {}", res.model_used).yellow()
                );
            }
            // Record budget for each failed attempt plus the successful one
            for _ in 0..res.attempts.len() {
                let _ = budget_tracker.record_call(false);
            }
            let _ = budget_tracker.record_call(res.tier_used == Tier::Paid);
        }
        Err(e) => {
            eprintln!("\n{} {}", "Error dispatching to persona:".red().bold(), e);
        }
    }

    Ok(())
}

#[allow(clippy::too_many_arguments)]
async fn run_pipeline(
    pipeline: &PipelineConfig,
    config: &AppConfig,
    budget_tracker: &mut BudgetTracker,
    model_registry: &ModelRegistry,
    client: &OpenRouterClient,
    prompt: &str,
    allow_stealth: bool,
    require_paid: bool,
) -> Result<(), Box<dyn std::error::Error>> {
    println!(
        "\n{} Starting pipeline '{}' ({} steps)",
        "⚡".bright_cyan(),
        pipeline.name.bold(),
        pipeline.steps.len()
    );

    let (event_tx, mut event_rx) = mpsc::channel::<PipelineEvent>(100);

    let exec_future = PipelineExecutor::execute(
        pipeline,
        &config.persona_registry,
        model_registry,
        client,
        prompt,
        allow_stealth,
        require_paid,
        Some(event_tx),
    );

    let event_handler = tokio::spawn(async move {
        let mut current_spinner: Option<ProgressBar> = None;

        while let Some(event) = event_rx.recv().await {
            match event {
                PipelineEvent::StepStarted {
                    step_index,
                    total_steps,
                    persona,
                    model,
                    ..
                } => {
                    if let Some(sp) = current_spinner.take() {
                        sp.finish_and_clear();
                    }
                    let badge = get_persona_color_tag(&persona);
                    println!(
                        "\n--- Step {}/{} {} ({}) ---",
                        step_index + 1,
                        total_steps,
                        badge,
                        model.dimmed()
                    );
                    let sp = create_spinner("Thinking...");
                    current_spinner = Some(sp);
                }
                PipelineEvent::Chunk { delta, .. } => {
                    if let Some(sp) = current_spinner.take() {
                        sp.finish_and_clear();
                    }
                    print!("{delta}");
                    let _ = io::stdout().flush();
                }
                PipelineEvent::StepCompleted { fell_back, .. } => {
                    if let Some(sp) = current_spinner.take() {
                        sp.finish_and_clear();
                    }
                    println!();
                    if fell_back {
                        eprintln!("{}", "⚠ Fallback occurred during step execution".yellow());
                    }
                }
            }
        }
    });

    let results = exec_future.await?;
    let _ = event_handler.await;

    for r in &results {
        for _ in 0..r.attempts.len() {
            let _ = budget_tracker.record_call(false);
        }
        let _ = budget_tracker.record_call(r.tier_used == Tier::Paid);
    }

    Ok(())
}