deepseek-rust-cli 1.20.6

use std::{
    io::{self},
    sync::Arc,
    time::{Duration, Instant},
};

use anyhow::Result;
use clap::{CommandFactory, Parser};
use crossterm::{
    cursor,
    event::{self, Event},
    execute, terminal,
};
use deepseek_rust_cli::{
    agent::{
        agent::DeepSeekAgent, commands::process_command, mentions::process_mentions,
        types::AgentEvent,
    },
    cli::{Args, ShellType},
    config::{get_api_key, init_workspace, load_config},
    logger::init_logger,
    tui::event_loop::{EventLoop, TuiEvent},
};
use tokio::sync::{mpsc, Mutex};

#[tokio::main]
async fn main() -> Result<()> {
    deepseek_rust_cli::tools::base::init_startup_dir();
    let args = Args::parse();

    // Handle --generate-completion early (no TUI needed)
    if let Some(shell) = &args.generate_completion {
        let mut cmd = Args::command();
        let name = cmd.get_name().to_string();
        let mut stdout = io::stdout();
        match shell {
            ShellType::Bash => {
                clap_complete::generate(clap_complete::shells::Bash, &mut cmd, name, &mut stdout)
            }
            ShellType::Zsh => {
                clap_complete::generate(clap_complete::shells::Zsh, &mut cmd, name, &mut stdout)
            }
            ShellType::Fish => {
                clap_complete::generate(clap_complete::shells::Fish, &mut cmd, name, &mut stdout)
            }
            ShellType::PowerShell => clap_complete::generate(
                clap_complete::shells::PowerShell,
                &mut cmd,
                name,
                &mut stdout,
            ),
            ShellType::Elvish => {
                clap_complete::generate(clap_complete::shells::Elvish, &mut cmd, name, &mut stdout)
            }
        }
        return Ok(());
    }

    let mut stdout = io::stdout();
    execute!(
        stdout,
        terminal::Clear(terminal::ClearType::All),
        cursor::MoveTo(0, 0)
    )?;

    init_workspace();

    let api_key = get_api_key()?;
    let mut config = load_config();
    init_logger(args.debug || config.debug);
    if args.danger_accept_invalid_certs {
        config.danger_accept_invalid_certs = true;
    }
    let mut agent = DeepSeekAgent::new(api_key, config, args.session);
    agent.auto_approve = args.auto_approve;

    let agent = Arc::new(Mutex::new(agent));

    // Extract tokens BEFORE spawning agent task to ensure they match
    let cancel_token = {
        let a = agent.try_lock().expect("agent should not be locked yet");
        a.cancel_token.clone()
    };
    let run_id = {
        let a = agent.try_lock().expect("agent should not be locked yet");
        a.run_id.clone()
    };

    let (tui_tx, tui_rx) = mpsc::channel(100);
    let (app_tx, mut app_rx) = mpsc::channel(1);
    let (cmd_tx, mut cmd_rx) = mpsc::channel::<(usize, String)>(100);

    // Input loop
    let tui_tx_for_input = tui_tx.clone();
    tokio::spawn(async move {
        let tick_rate = Duration::from_millis(100);
        let mut last_tick = Instant::now();
        loop {
            let timeout = tick_rate
                .checked_sub(last_tick.elapsed())
                .unwrap_or(Duration::from_secs(0));

            if event::poll(timeout).unwrap_or(false) {
                match event::read() {
                    Ok(Event::Key(key)) => {
                        let _ = tui_tx_for_input.send(TuiEvent::Input(key)).await;
                    }
                    Ok(Event::Mouse(mouse)) => {
                        let _ = tui_tx_for_input.send(TuiEvent::Mouse(mouse)).await;
                    }
                    Ok(Event::Paste(text)) => {
                        let _ = tui_tx_for_input.send(TuiEvent::Paste(text)).await;
                    }
                    Ok(Event::Resize(width, height)) => {
                        let _ = tui_tx_for_input.send(TuiEvent::Resize(width, height)).await;
                    }
                    Ok(_) => {}
                    Err(e) => {
                        tracing::warn!("Input event error: {}", e);
                    }
                }
            }
            if last_tick.elapsed() >= tick_rate {
                let _ = tui_tx_for_input.send(TuiEvent::Tick).await;
                last_tick = Instant::now();
            }
        }
    });

    // Agent processing task
    let agent_clone = agent.clone();
    let tui_tx_for_agent = tui_tx.clone();

    tokio::spawn(async move {
        tracing::info!("Agent task started, waiting for commands...");
        while let Some((cmd_run_id, cmd)) = cmd_rx.recv().await {
            tracing::info!("Agent received command: {}", cmd);
            let mut agent_lock = agent_clone.lock().await;
            tracing::info!("Agent lock acquired for command: {}", cmd);

            // If the command is from an aborted session, skip it completely.
            // This clears the queue of old commands.
            if cmd_run_id != agent_lock.run_id.load(std::sync::atomic::Ordering::SeqCst) {
                tracing::warn!("Skipping stale command (run_id mismatch): {}", cmd);
                continue;
            }

            let (agent_event_tx, mut agent_event_rx) = mpsc::channel(100);

            let tui_tx_inner = tui_tx_for_agent.clone();
            tokio::spawn(async move {
                while let Some(ev) = agent_event_rx.recv().await {
                    let _ = tui_tx_inner.send(TuiEvent::Agent(ev)).await;
                }
            });

            // Handle slash commands
            if cmd.starts_with('/') {
                tracing::info!("Processing slash command: {}", cmd);
                match process_command(&mut agent_lock, &cmd).await {
                    Ok(Some(response)) => {
                        let tu = agent_lock.token_usage.clone();
                        let _ = agent_event_tx
                            .send(AgentEvent::Content { content: response })
                            .await;
                        let _ = agent_event_tx
                            .send(AgentEvent::Done { token_usage: tu })
                            .await;
                        continue;
                    }
                    Ok(None) => {
                        // Not a recognized command, proceed to chat
                        tracing::info!("Slash command not recognized, trying as chat: {}", cmd);
                    }
                    Err(e) => {
                        let tu = agent_lock.token_usage.clone();
                        let _ = agent_event_tx
                            .send(AgentEvent::Error {
                                content: format!("Command error: {}", e),
                            })
                            .await;
                        let _ = agent_event_tx
                            .send(AgentEvent::Done { token_usage: tu })
                            .await;
                        continue;
                    }
                }
            }

            let processed_cmd = process_mentions(&cmd);
            tracing::info!("Calling chat_stream for: {}", processed_cmd);
            let chat_result = agent_lock
                .chat_stream(processed_cmd, agent_event_tx.clone(), &mut app_rx)
                .await;
            match &chat_result {
                Ok(()) => tracing::info!("chat_stream completed successfully"),
                Err(e) => tracing::error!("chat_stream failed: {}", e),
            }
            // Only send Done if not aborted (Aborted event already sent by chat_stream)
            if !agent_lock
                .cancel_token
                .lock()
                .unwrap_or_else(|e| e.into_inner())
                .is_cancelled()
            {
                let tu = agent_lock.token_usage.clone();
                let _ = agent_event_tx
                    .send(AgentEvent::Done { token_usage: tu })
                    .await;
            }
            agent_lock.reset_cancel();
            tracing::info!("Agent done processing command: {}", cmd);
        }
        tracing::warn!("Agent task cmd_rx closed, exiting");
    });

    // Start TUI
    let event_loop = EventLoop::new(
        tui_rx,
        tui_tx.clone(),
        app_tx,
        cmd_tx,
        agent.clone(),
        cancel_token,
        run_id,
    );

    let res = tokio::select! {
        r = event_loop.run() => r,
        _ = tokio::signal::ctrl_c() => {
            tracing::warn!("Ctrl+C signal received, shutting down gracefully");
            Ok(String::new())
        }
    };

    execute!(
        io::stdout(),
        terminal::Clear(terminal::ClearType::All),
        cursor::MoveTo(0, 0)
    )?;

    if let Err(e) = res {
        println!("\n❌ UI error: {}", e);
        std::process::exit(1);
    }

    std::process::exit(0);
}