zilliz 1.4.3

use std::time::Instant;

use crossterm::event::{KeyCode, KeyEvent, KeyModifiers};

use crate::auth::device_code;

use super::app::{App, Screen, WizardMsg};
use super::wizard::{AuthMethod, Region as WizardRegion, WizardFocus, WizardState};

/// Handle a key event. Routed by the top of the screen stack.
pub fn handle_key(app: &mut App, key: KeyEvent) {
    // Help overlay: esc / ? closes it.
    if *app.current_screen() == Screen::Help {
        if matches!(key.code, KeyCode::Esc | KeyCode::Char('?')) {
            app.screen_stack.pop();
        } else if matches!(
            (key.code, key.modifiers),
            (KeyCode::Char('c'), m) if m.contains(KeyModifiers::CONTROL)
        ) {
            app.should_quit = true;
        }
        return;
    }

    // Global Ctrl+C quits from any screen.
    if matches!(key.code, KeyCode::Char('c')) && key.modifiers.contains(KeyModifiers::CONTROL) {
        cancel_wizard(app);
        app.should_quit = true;
        return;
    }

    match app.current_screen().clone() {
        Screen::Home => handle_home(app, key),
        Screen::SignInRegion => handle_region(app, key),
        Screen::SignInMethod => handle_method(app, key),
        Screen::SignInBrowser => handle_browser(app, key),
        Screen::SignInApiKey => handle_api_key(app, key),
        Screen::LogoutConfirm => handle_logout_confirm(app, key),
        Screen::Help => {}
    }
}

// --- Home ---------------------------------------------------------------

fn handle_home(app: &mut App, key: KeyEvent) {
    match key.code {
        KeyCode::Char('q') | KeyCode::Esc => app.should_quit = true,
        KeyCode::Char('?') => app.screen_stack.push(Screen::Help),
        KeyCode::Char('l') => {
            if app.auth.is_signed_in() {
                app.screen_stack.push(Screen::LogoutConfirm);
            } else {
                enter_wizard(app);
            }
        }
        _ => {}
    }
}

fn handle_logout_confirm(app: &mut App, key: KeyEvent) {
    match key.code {
        KeyCode::Char('y') | KeyCode::Char('Y') | KeyCode::Enter => {
            // Mirror `cli/auth.rs::logout`: clear credentials, context, and
            // any persisted control-plane endpoint override.
            let _ = app.config_mgr.clear_login_data();
            let _ = app.config_mgr.clear_context();
            let _ = app.config_mgr.clear_control_plane_endpoint();
            app.refresh_auth();
            // Sign-out resets the home counts so a future sign-in re-fetches.
            app.abort_home_counts_fetch();
            app.screen_stack = vec![Screen::Home];
        }
        KeyCode::Char('n') | KeyCode::Char('N') | KeyCode::Esc => {
            app.screen_stack.pop();
        }
        KeyCode::Char('?') => {
            app.screen_stack.push(Screen::Help);
        }
        _ => {}
    }
}

// --- Region (step 1) -----------------------------------------------------

fn handle_region(app: &mut App, key: KeyEvent) {
    let Some(w) = app.wizard.as_mut() else {
        cancel_wizard(app);
        return;
    };
    match key.code {
        KeyCode::Esc => {
            cancel_wizard(app);
        }
        KeyCode::Char('?') => {
            app.screen_stack.push(Screen::Help);
        }
        KeyCode::Up | KeyCode::Char('k') => {
            w.region_cursor = if w.region_cursor == 0 { 1 } else { 0 };
        }
        KeyCode::Down | KeyCode::Char('j') => {
            w.region_cursor = (w.region_cursor + 1) % 2;
        }
        KeyCode::Char('g') => {
            w.region_cursor = 0;
            w.region = Some(WizardRegion::Global);
            app.screen_stack.push(Screen::SignInMethod);
        }
        KeyCode::Char('c') => {
            w.region_cursor = 1;
            w.region = Some(WizardRegion::China);
            // China region is API-key only — skip method step.
            app.screen_stack.push(Screen::SignInApiKey);
        }
        KeyCode::Enter => {
            let region = w.region_after_cursor();
            w.region = Some(region);
            if region.supports_browser() {
                app.screen_stack.push(Screen::SignInMethod);
            } else {
                app.screen_stack.push(Screen::SignInApiKey);
            }
        }
        _ => {}
    }
}

// --- Method (step 2) -----------------------------------------------------

fn handle_method(app: &mut App, key: KeyEvent) {
    let Some(w) = app.wizard.as_mut() else {
        cancel_wizard(app);
        return;
    };
    match key.code {
        KeyCode::Esc => cancel_wizard(app),
        KeyCode::Char('?') => app.screen_stack.push(Screen::Help),
        KeyCode::Char('b') => {
            app.screen_stack.pop();
        }
        KeyCode::Up | KeyCode::Char('k') => {
            w.method_cursor = if w.method_cursor == 0 { 1 } else { 0 };
        }
        KeyCode::Down | KeyCode::Char('j') => {
            w.method_cursor = (w.method_cursor + 1) % 2;
        }
        KeyCode::Char('1') => {
            w.method_cursor = 0;
            w.method = Some(AuthMethod::Browser);
            start_browser_flow(app);
        }
        KeyCode::Char('2') => {
            w.method_cursor = 1;
            w.method = Some(AuthMethod::ApiKey);
            app.screen_stack.push(Screen::SignInApiKey);
        }
        KeyCode::Enter => {
            let method = w.method_after_cursor();
            w.method = Some(method);
            match method {
                AuthMethod::Browser => start_browser_flow(app),
                AuthMethod::ApiKey => app.screen_stack.push(Screen::SignInApiKey),
            }
        }
        _ => {}
    }
}

// --- Browser (step 3a) ---------------------------------------------------

fn handle_browser(app: &mut App, key: KeyEvent) {
    match key.code {
        KeyCode::Esc => cancel_wizard(app),
        KeyCode::Char('?') => app.screen_stack.push(Screen::Help),
        KeyCode::Char('o') => {
            if let Some(url) = app
                .wizard
                .as_ref()
                .and_then(|w| w.device_code.as_ref())
                .map(|s| s.response.verification_uri_complete.clone())
            {
                let _ = device_code::open_browser(&url);
            }
        }
        _ => {}
    }
}

// --- API key (step 3b) ---------------------------------------------------

fn handle_api_key(app: &mut App, key: KeyEvent) {
    let Some(w) = app.wizard.as_mut() else {
        cancel_wizard(app);
        return;
    };
    match (key.code, key.modifiers) {
        (KeyCode::Esc, _) => {
            cancel_wizard(app);
        }
        (KeyCode::Char('?'), m)
            if !m.contains(KeyModifiers::CONTROL) && w.api_key_focus != WizardFocus::Input =>
        {
            app.screen_stack.push(Screen::Help);
        }
        (KeyCode::Char('h'), m) if m.contains(KeyModifiers::CONTROL) => {
            w.api_key_visible = !w.api_key_visible;
        }
        (KeyCode::Char('v'), m) if m.contains(KeyModifiers::CONTROL) => {
            // Clipboard access is best-effort; we don't depend on a clipboard
            // crate in this change. Surface a non-fatal hint instead of failing.
            w.error = Some("Clipboard unavailable — paste with terminal shortcut.".to_string());
        }
        (KeyCode::Char('b'), m)
            if !m.contains(KeyModifiers::CONTROL) && w.api_key_focus != WizardFocus::Input =>
        {
            app.screen_stack.pop();
        }
        (KeyCode::Tab, _) | (KeyCode::Down, _) | (KeyCode::Up, _) => {
            w.api_key_focus = match w.api_key_focus {
                WizardFocus::Input => WizardFocus::Save,
                WizardFocus::Save => WizardFocus::Input,
            };
        }
        (KeyCode::Char(' '), _)
            if w.api_key_focus == WizardFocus::Input && w.api_key_buf.chars().count() < 256 =>
        {
            w.api_key_buf.push(' ');
        }
        (KeyCode::Backspace, _) if w.api_key_focus == WizardFocus::Input => {
            w.api_key_buf.pop();
        }
        (KeyCode::Char(c), m)
            if w.api_key_focus == WizardFocus::Input
                && !m.contains(KeyModifiers::CONTROL)
                && !m.contains(KeyModifiers::ALT)
                && w.api_key_buf.chars().count() < 256 =>
        {
            w.api_key_buf.push(c);
        }
        (KeyCode::Enter, _) => {
            let buf = w.api_key_buf.trim().to_string();
            if buf.is_empty() {
                w.error = Some("Enter your API key, or press esc to cancel.".to_string());
                return;
            }
            let endpoint = w
                .region
                .map(|r| r.endpoint())
                .unwrap_or("https://api.cloud.zilliz.com");
            // Persist the endpoint *before* the key so a failure can't leave a
            // key stored against the wrong region (e.g. a China key resolving
            // against the Global endpoint, which would fail confusingly later).
            if let Err(e) = persist_endpoint(&app.config_mgr, endpoint) {
                w.error = Some(format!("Failed to set endpoint: {}", e));
                return;
            }
            if let Err(e) = device_code::save_api_key(&app.config_mgr, &buf) {
                w.error = Some(format!("Failed to save: {}", e));
                return;
            }
            finish_sign_in(app);
        }
        _ => {}
    }
}

// --- Transitions ---------------------------------------------------------

fn enter_wizard(app: &mut App) {
    app.wizard = Some(WizardState::new());
    app.screen_stack.push(Screen::SignInRegion);
}

pub fn cancel_wizard(app: &mut App) {
    if let Some(w) = app.wizard.as_ref() {
        w.cancel.cancel();
    }
    app.wizard = None;
    app.screen_stack = vec![Screen::Home];
}

fn finish_sign_in(app: &mut App) {
    if let Some(w) = app.wizard.as_ref() {
        w.cancel.cancel();
    }
    app.wizard = None;
    app.screen_stack = vec![Screen::Home];
    app.refresh_auth();
    if app.auth.is_signed_in() {
        app.spawn_home_counts_fetch();
    }
}

fn start_browser_flow(app: &mut App) {
    // Auth0 config is required; if missing we fall back to API-key entry
    // with an inline error so the user is not stranded.
    let auth_config = match app.models.control_plane.auth.clone() {
        Some(cfg) if !cfg.auth0_domain.is_empty() => cfg,
        _ => {
            if let Some(w) = app.wizard.as_mut() {
                w.error = Some("Browser sign-in is unavailable in this build.".to_string());
            }
            app.screen_stack.push(Screen::SignInApiKey);
            return;
        }
    };

    let tx = app.msg_tx.clone();
    let cancel = app
        .wizard
        .as_ref()
        .map(|w| w.cancel.clone())
        .unwrap_or_default();

    // The device-code flow needs a Tokio runtime. Handler unit tests run
    // without one (and future non-runtime callers would otherwise panic in
    // `tokio::spawn`), so guard like `spawn_home_counts_fetch` and surface an
    // inline error rather than crashing.
    let Ok(handle) = tokio::runtime::Handle::try_current() else {
        if let Some(w) = app.wizard.as_mut() {
            w.error = Some("Browser sign-in needs an async runtime.".to_string());
        }
        app.screen_stack.push(Screen::SignInApiKey);
        return;
    };

    app.screen_stack.push(Screen::SignInBrowser);

    handle.spawn(async move {
        match device_code::request_device_code(&auth_config).await {
            Ok(resp) => {
                // Best-effort browser open before we even ship the code into UI.
                let _ = device_code::open_browser(&resp.verification_uri_complete);
                let interval = resp.interval;
                let expires_in = resp.expires_in;
                let device = resp.device_code.clone();
                let _ = tx.send(WizardMsg::DeviceCodeReady(resp));

                match device_code::poll_for_token(
                    &auth_config,
                    &device,
                    interval,
                    expires_in,
                    cancel.clone(),
                )
                .await
                {
                    Ok(token) => {
                        // Bail out if the wizard was cancelled (or superseded by
                        // a fresh flow, which cancels this token) while polling,
                        // so a stale poller can't silently sign the user in.
                        if cancel.is_cancelled() {
                            return;
                        }
                        // Exchange Auth0 token → CLI credentials, same as
                        // `zilliz login` does after browser auth completes.
                        match device_code::exchange_token(&auth_config, &token.access_token).await {
                            Ok(payload) => {
                                if cancel.is_cancelled() {
                                    return;
                                }
                                let _ = tx.send(WizardMsg::LoginExchanged(payload));
                            }
                            Err(e) => {
                                let _ = tx.send(WizardMsg::ExchangeError(e.to_string()));
                            }
                        }
                    }
                    Err(e) => {
                        let _ = tx.send(WizardMsg::PollError(e.to_string()));
                    }
                }
            }
            Err(e) => {
                let _ = tx.send(WizardMsg::DeviceCodeError(e.to_string()));
            }
        }
    });
}

pub fn handle_wizard_msg(app: &mut App, msg: WizardMsg) {
    match msg {
        WizardMsg::DeviceCodeReady(resp) => {
            if let Some(w) = app.wizard.as_mut() {
                w.device_code = Some(super::wizard::DeviceCodeSession {
                    response: resp,
                    started_at: Instant::now(),
                });
                w.error = None;
            }
        }
        WizardMsg::DeviceCodeError(err) => {
            if let Some(w) = app.wizard.as_mut() {
                w.error = Some(err);
            }
        }
        WizardMsg::LoginExchanged(payload) => {
            // Ignore a late success from a cancelled or superseded flow: once
            // the wizard is gone (cancel/finish set it to `None`) we must not
            // write credentials, or a stale Auth0 poller could silently sign
            // the user in.
            if app.wizard.is_none() {
                return;
            }
            // Browser flow always targets Global; clear any persisted override
            // *before* saving credentials so we can never end up with a Global
            // Auth0 session still pointed at a previous CN/dev endpoint.
            if let Err(e) = app.config_mgr.clear_control_plane_endpoint() {
                if let Some(w) = app.wizard.as_mut() {
                    w.error = Some(format!("Failed to reset endpoint: {}", e));
                }
                return;
            }
            // Persist the same way `cli/auth.rs::login_with_browser` does.
            if let Err(e) = app.config_mgr.save_login_data(
                &payload.user_id,
                &payload.email,
                &payload.name,
                &payload.orgs,
            ) {
                if let Some(w) = app.wizard.as_mut() {
                    w.error = Some(format!("Failed to save credentials: {}", e));
                }
                return;
            }
            finish_sign_in(app);
        }
        WizardMsg::PollError(err) | WizardMsg::ExchangeError(err) => {
            if let Some(w) = app.wizard.as_mut() {
                w.error = Some(err);
            }
        }
    }
}

fn persist_endpoint(
    cfg: &crate::config::manager::ConfigManager,
    endpoint: &str,
) -> anyhow::Result<()> {
    // Only persist non-default endpoints (China / dev) so we don't write a
    // useless entry for Global.
    if endpoint == "https://api.cloud.zilliz.com" {
        cfg.clear_control_plane_endpoint()?;
    } else {
        cfg.set_control_plane_endpoint(endpoint)?;
    }
    Ok(())
}