pub mod app;
pub mod ui;

use app::{AppState, TuiMode};
use brokecrew::budget::BudgetTracker;
use brokecrew::client::OpenRouterClient;
use brokecrew::config::AppConfig;
use brokecrew::persona::Tier;
use brokecrew::pipeline::{PipelineEvent, PipelineExecutor};
use brokecrew::registry::ModelRegistry;
use brokecrew::router::{RouteStage, Router};
use crossterm::{
    event::{Event, EventStream, KeyCode, KeyEvent, KeyModifiers},
    execute,
    terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen},
};
use futures_util::StreamExt;
use ratatui::{backend::CrosstermBackend, Terminal};
use std::io;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::mpsc;

struct TerminalGuard;

impl Drop for TerminalGuard {
    fn drop(&mut self) {
        let _ = disable_raw_mode();
        let _ = execute!(io::stdout(), LeaveAlternateScreen, crossterm::cursor::Show);
    }
}

pub enum WorkerEvent {
    RouterResolved {
        persona: String,
        stage: RouteStage,
    },
    StepStarted {
        step_index: usize,
        total_steps: usize,
        persona: String,
        model: String,
        tier: Tier,
    },
    Chunk(String),
    StepCompleted {
        fell_back: bool,
    },
    Finished {
        fell_back: bool,
        model_used: String,
        error: Option<String>,
        calls_made: usize,
        paid_calls_made: usize,
    },
}

pub async fn run_tui(
    config: AppConfig,
    mut budget_tracker: BudgetTracker,
    model_registry: Arc<ModelRegistry>,
    allow_stealth: bool,
    escalate: bool,
) -> Result<(), Box<dyn std::error::Error>> {
    // Enable raw mode and alternate screen
    enable_raw_mode()?;
    let mut stdout = io::stdout();
    execute!(stdout, EnterAlternateScreen, crossterm::cursor::Show)?;
    let _guard = TerminalGuard;

    let backend = CrosstermBackend::new(stdout);
    let mut terminal = Terminal::new(backend)?;

    // Prepare personas list for AppState
    let personas_summary: Vec<(String, String)> = config
        .persona_registry
        .list()
        .iter()
        .map(|p| {
            let desc = if !p.trigger_words.is_empty() {
                format!("triggers: {}", p.trigger_words.join(", "))
            } else {
                p.system_prompt.chars().take(50).collect::<String>()
            };
            (p.name.clone(), desc)
        })
        .collect();

    let (minute_limit, daily_limit) = match budget_tracker.check_budget() {
        Ok(s) => (s.minute_limit, s.daily_limit),
        Err(_) => (20, 50),
    };

    let mut app = AppState::new(
        personas_summary,
        config.pipelines.clone(),
        minute_limit,
        daily_limit,
    );
    app.allow_stealth = allow_stealth || config.allow_stealth;
    app.escalate = escalate;

    // Refresh model catalog in background
    let reg_clone = Arc::clone(&model_registry);
    let url_clone = config.openrouter_base_url.clone();
    tokio::spawn(async move {
        let _ = reg_clone.fetch_live(&url_clone).await;
    });

    let client = Arc::new(OpenRouterClient::with_base_url(
        config.api_key.clone().unwrap_or_default(),
        config.openrouter_base_url.clone(),
    ));
    let router = Arc::new(Router::new(
        config.default_persona.clone(),
        config.router_model.clone(),
        config.complexity_threshold,
    ));

    let (worker_tx, mut worker_rx) = mpsc::channel::<WorkerEvent>(100);

    let mut event_stream = EventStream::new();
    let mut tick_interval = tokio::time::interval(Duration::from_millis(80));

    loop {
        // Update budget counters in AppState
        if let Ok(st) = budget_tracker.check_budget() {
            app.minute_requests = st.minute_requests;
            app.daily_requests = st.daily_requests;
            app.minute_limit = st.minute_limit;
            app.daily_limit = st.daily_limit;
        }
        app.models_cached = model_registry.model_count().await;

        // Render UI
        terminal.draw(|f| ui::render(f, &mut app))?;

        tokio::select! {
            _ = tick_interval.tick() => {
                if app.is_generating {
                    app.spinner_tick = app.spinner_tick.wrapping_add(1);
                }
            }

            Some(event) = worker_rx.recv() => {
                handle_worker_event(&mut app, &mut budget_tracker, event);
            }

            maybe_event = event_stream.next() => {
                match maybe_event {
                    Some(Ok(Event::Key(key_event))) => {
                        let should_quit = handle_key_event(
                            &mut app,
                            key_event,
                            &config,
                            &model_registry,
                            &client,
                            &router,
                            worker_tx.clone(),
                        );
                        if should_quit {
                            break;
                        }
                    }
                    Some(Ok(Event::Resize(_, _))) => {
                        terminal.autoresize()?;
                    }
                    _ => {}
                }
            }
        }
    }

    Ok(())
}

fn handle_worker_event(app: &mut AppState, budget_tracker: &mut BudgetTracker, event: WorkerEvent) {
    match event {
        WorkerEvent::RouterResolved { persona, stage } => {
            let stage_info = match stage {
                RouteStage::StageAKeyword => "Stage A keyword match",
                RouteStage::StageBLlm => "Stage B LLM classification",
            };
            app.add_system_message(format!("⚡ Routed prompt to [{persona}] via {stage_info}"));
        }

        WorkerEvent::StepStarted {
            step_index,
            total_steps,
            persona,
            model,
            tier,
        } => {
            if total_steps > 1 {
                app.start_pipeline_step_message(step_index, total_steps, &persona, &model, tier);
            } else {
                app.start_assistant_message(&persona, &model, tier);
            }
        }

        WorkerEvent::Chunk(chunk) => {
            app.append_to_current_message(&chunk);
        }

        WorkerEvent::StepCompleted { fell_back } => {
            if fell_back {
                app.mark_current_message_fallback();
            }
        }

        WorkerEvent::Finished {
            fell_back,
            model_used,
            error,
            calls_made,
            paid_calls_made,
        } => {
            app.is_generating = false;
            app.status_message = format!("Ready ({model_used})");
            if fell_back {
                app.mark_current_message_fallback();
            }

            if let Some(err) = error {
                app.add_system_message(format!("❌ Error: {err}"));
            }

            for _ in 0..calls_made {
                let _ = budget_tracker.record_call(false);
            }
            for _ in 0..paid_calls_made {
                let _ = budget_tracker.record_call(true);
            }
        }
    }
}

fn handle_key_event(
    app: &mut AppState,
    key: KeyEvent,
    config: &AppConfig,
    model_registry: &Arc<ModelRegistry>,
    client: &Arc<OpenRouterClient>,
    router: &Arc<Router>,
    worker_tx: mpsc::Sender<WorkerEvent>,
) -> bool {
    // Global quit on Ctrl+C
    if key.modifiers.contains(KeyModifiers::CONTROL) && key.code == KeyCode::Char('c') {
        return true;
    }

    match app.mode {
        TuiMode::PersonaSelect => match key.code {
            KeyCode::Esc => {
                app.mode = TuiMode::Normal;
            }
            KeyCode::Up => {
                app.selected_persona_idx = app.selected_persona_idx.saturating_sub(1);
            }
            KeyCode::Down => {
                if app.selected_persona_idx + 1 < app.persona_items.len() {
                    app.selected_persona_idx += 1;
                }
            }
            KeyCode::Enter => {
                app.active_pipeline_for_prompt = None;
                app.mode = TuiMode::Normal;
            }
            _ => {}
        },

        TuiMode::PipelineSelect => match key.code {
            KeyCode::Esc => {
                app.mode = TuiMode::Normal;
            }
            KeyCode::Up => {
                app.selected_pipeline_idx = app.selected_pipeline_idx.saturating_sub(1);
            }
            KeyCode::Down => {
                if app.selected_pipeline_idx + 1 < app.pipelines.len() {
                    app.selected_pipeline_idx += 1;
                }
            }
            KeyCode::Enter => {
                if let Some(pipe) = app.pipelines.get(app.selected_pipeline_idx) {
                    app.active_pipeline_for_prompt = Some(pipe.name.clone());
                }
                app.mode = TuiMode::Normal;
            }
            _ => {}
        },

        TuiMode::Normal => {
            // Mode switching and shortcuts
            if key.modifiers.contains(KeyModifiers::CONTROL) {
                match key.code {
                    KeyCode::Char('p') => {
                        app.mode = TuiMode::PipelineSelect;
                        return false;
                    }
                    KeyCode::Char('l') => {
                        app.clear_chat();
                        return false;
                    }
                    _ => {}
                }
            }

            match key.code {
                KeyCode::Esc => {
                    return true;
                }
                KeyCode::Tab => {
                    app.mode = TuiMode::PersonaSelect;
                }
                KeyCode::PageUp => {
                    app.scroll_up(5);
                }
                KeyCode::PageDown => {
                    app.scroll_down(5);
                }
                KeyCode::Left => {
                    app.move_cursor_left();
                }
                KeyCode::Right => {
                    app.move_cursor_right();
                }
                KeyCode::Home => {
                    app.move_cursor_start();
                }
                KeyCode::End => {
                    app.move_cursor_end();
                }
                KeyCode::Backspace => {
                    app.delete_char_before_cursor();
                }
                KeyCode::Delete => {
                    app.delete_char_under_cursor();
                }
                KeyCode::Enter => {
                    let prompt = app.input.trim().to_string();
                    if !prompt.is_empty() && !app.is_generating {
                        app.add_user_message(prompt.clone());
                        app.input.clear();
                        app.cursor_pos = 0;
                        app.is_generating = true;
                        app.status_message = "Dispatching...".to_string();

                        spawn_dispatch_task(
                            prompt,
                            app,
                            config,
                            model_registry,
                            client,
                            router,
                            worker_tx,
                        );
                    }
                }
                KeyCode::Char(c) => {
                    app.insert_char(c);
                }
                _ => {}
            }
        }
    }

    false
}

fn spawn_dispatch_task(
    prompt: String,
    app: &AppState,
    config: &AppConfig,
    model_registry: &Arc<ModelRegistry>,
    client: &Arc<OpenRouterClient>,
    router: &Arc<Router>,
    worker_tx: mpsc::Sender<WorkerEvent>,
) {
    let pipeline_name = app.active_pipeline_for_prompt.clone();
    let is_auto = app.is_auto_route();
    let direct_persona = app.current_persona_name().to_string();
    let allow_stealth = app.allow_stealth;
    let escalate = app.escalate;

    let persona_registry = config.persona_registry.clone();
    let pipelines = config.pipelines.clone();
    let default_persona = config.default_persona.clone();
    let model_registry = Arc::clone(model_registry);
    let client = Arc::clone(client);
    let router = Arc::clone(router);

    tokio::spawn(async move {
        if let Some(pipe_name) = pipeline_name {
            if let Some(pipe) = pipelines.iter().find(|p| p.name == pipe_name) {
                let (event_tx, mut event_rx) = mpsc::channel::<PipelineEvent>(50);
                let w_tx = worker_tx.clone();

                let forwarder = tokio::spawn(async move {
                    while let Some(event) = event_rx.recv().await {
                        match event {
                            PipelineEvent::StepStarted {
                                step_index,
                                total_steps,
                                persona,
                                model,
                                tier,
                            } => {
                                let _ = w_tx
                                    .send(WorkerEvent::StepStarted {
                                        step_index,
                                        total_steps,
                                        persona,
                                        model,
                                        tier,
                                    })
                                    .await;
                            }
                            PipelineEvent::Chunk { delta, .. } => {
                                let _ = w_tx.send(WorkerEvent::Chunk(delta)).await;
                            }
                            PipelineEvent::StepCompleted { fell_back, .. } => {
                                let _ = w_tx.send(WorkerEvent::StepCompleted { fell_back }).await;
                            }
                        }
                    }
                });

                let res = PipelineExecutor::execute(
                    pipe,
                    &persona_registry,
                    &model_registry,
                    client.as_ref(),
                    &prompt,
                    allow_stealth,
                    escalate,
                    Some(event_tx),
                )
                .await;

                let _ = forwarder.await;

                match res {
                    Ok(results) => {
                        let total_calls: usize = results.iter().map(|r| r.attempts.len()).sum();
                        let paid_calls: usize =
                            results.iter().filter(|r| r.tier_used == Tier::Paid).count();
                        let fell_back = results.iter().any(|r| r.fell_back);
                        let _ = worker_tx
                            .send(WorkerEvent::Finished {
                                fell_back,
                                model_used: "pipeline".to_string(),
                                error: None,
                                calls_made: total_calls + results.len(),
                                paid_calls_made: paid_calls,
                            })
                            .await;
                    }
                    Err(e) => {
                        let _ = worker_tx
                            .send(WorkerEvent::Finished {
                                fell_back: false,
                                model_used: "none".to_string(),
                                error: Some(e.to_string()),
                                calls_made: 1,
                                paid_calls_made: 0,
                            })
                            .await;
                    }
                }
                return;
            }
        }

        // Determine target persona
        let (target_persona, route_stage) = if is_auto {
            match router
                .route(&persona_registry, client.as_ref(), &prompt, escalate)
                .await
            {
                Ok(decision) => (decision.persona, Some(decision.stage_used)),
                Err(_) => (default_persona, None),
            }
        } else {
            (direct_persona, None)
        };

        if let Some(stage) = route_stage {
            let _ = worker_tx
                .send(WorkerEvent::RouterResolved {
                    persona: target_persona.clone(),
                    stage,
                })
                .await;
        }

        // Discover or dispatch persona
        let (chunk_tx, mut chunk_rx) = mpsc::channel::<String>(100);
        let w_tx = worker_tx.clone();

        // Notify UI of step started
        let _ = worker_tx
            .send(WorkerEvent::StepStarted {
                step_index: 0,
                total_steps: 1,
                persona: target_persona.clone(),
                model: "evaluating candidates...".to_string(),
                tier: Tier::Free,
            })
            .await;

        let forwarder = tokio::spawn(async move {
            while let Some(chunk) = chunk_rx.recv().await {
                let _ = w_tx.send(WorkerEvent::Chunk(chunk)).await;
            }
        });

        let res = PipelineExecutor::dispatch_persona(
            &target_persona,
            &persona_registry,
            &model_registry,
            client.as_ref(),
            &prompt,
            allow_stealth,
            escalate,
            Some(chunk_tx),
        )
        .await;

        let _ = forwarder.await;

        match res {
            Ok(disp) => {
                let _ = worker_tx
                    .send(WorkerEvent::Finished {
                        fell_back: disp.fell_back,
                        model_used: disp.model_used,
                        error: None,
                        calls_made: disp.attempts.len() + 1,
                        paid_calls_made: if disp.tier_used == Tier::Paid { 1 } else { 0 },
                    })
                    .await;
            }
            Err(e) => {
                let _ = worker_tx
                    .send(WorkerEvent::Finished {
                        fell_back: false,
                        model_used: "none".to_string(),
                        error: Some(e.to_string()),
                        calls_made: 1,
                        paid_calls_made: 0,
                    })
                    .await;
            }
        }
    });
}