vex-tui 0.6.2

use std::collections::{HashMap, HashSet};
use std::io;
use std::sync::Arc;
use std::time::Duration;

use crate::terminal_widget::VtTerminal;
use crate::ui;
use anyhow::Result;
use crossterm::event::{self, Event, KeyCode, KeyEvent, KeyModifiers};
use crossterm::execute;
use crossterm::terminal::{
    EnterAlternateScreen, LeaveAlternateScreen, disable_raw_mode, enable_raw_mode,
};
use ratatui::Terminal;
use ratatui::backend::CrosstermBackend;
use tokio::sync::{broadcast, mpsc, watch};
use uuid::Uuid;
use vex_hub::{FrontendCommand, FrontendEvent, Hub, HubState};

// ── Sidebar model ──────────────────────────────────────────────

/// A single entry in the flat sidebar list.
#[derive(Debug, Clone)]
pub(crate) enum SidebarEntry {
    // Top section: repos & workstreams with nested agents/shells
    RepoHeader {
        name: String,
    },
    WorkstreamItem {
        #[allow(dead_code)]
        repo: String,
        name: String,
    },
    NestedAgent {
        id: Uuid,
        label: String,
        needs: bool,
    },
    #[allow(dead_code)]
    NestedShell {
        id: Uuid,
        label: String,
    },
    // Divider between top and bottom sections
    Divider,
    // Bottom: agents section
    AgentSectionHeader,
    AgentItem {
        id: Uuid,
        label: String,
        needs: bool,
    },
    CreateAgentPlaceholder,
    // Bottom: shells section (non-agent shells only)
    ShellSectionHeader,
    ShellItem {
        id: Uuid,
        label: String,
    },
}

impl SidebarEntry {
    fn is_selectable(&self) -> bool {
        matches!(
            self,
            SidebarEntry::WorkstreamItem { .. }
                | SidebarEntry::NestedAgent { .. }
                | SidebarEntry::NestedShell { .. }
                | SidebarEntry::AgentItem { .. }
                | SidebarEntry::CreateAgentPlaceholder
                | SidebarEntry::ShellItem { .. }
        )
    }

    #[allow(dead_code)]
    fn is_in_agent_section(&self) -> bool {
        matches!(
            self,
            SidebarEntry::AgentSectionHeader
                | SidebarEntry::AgentItem { .. }
                | SidebarEntry::CreateAgentPlaceholder
        )
    }
}

// ── Right pane ─────────────────────────────────────────────────

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum RightPane {
    Empty,
    AgentConversation(Uuid),
    ShellTerminal(Uuid),
}

// ── Dialog screens (overlays) ──────────────────────────────────

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum Dialog {
    None,
    SpawnAgent {
        repo_idx: usize,
    },
    PromptAgent {
        shell_id: Uuid,
    },
    NewRepo {
        path_buf: String,
        name_buf: String,
        editing_name: bool,
    },
    NewWorkstream {
        repo_idx: usize,
        name_buf: String,
    },
}

// ── App state ──────────────────────────────────────────────────

pub struct App {
    pub state: HubState,
    pub cursor: usize,
    pub right_pane: RightPane,
    pub dialog: Dialog,
    pub running: bool,
    pub status_message: Option<String>,
    pub vt_terminals: HashMap<Uuid, VtTerminal>,
    pub agent_lines: HashMap<Uuid, Vec<String>>,
    pub prompt_buf: String,
    pub port: u16,
}

impl App {
    fn new(port: u16) -> Self {
        Self {
            state: HubState::default(),
            cursor: 0,
            right_pane: RightPane::Empty,
            dialog: Dialog::None,
            running: true,
            status_message: None,
            vt_terminals: HashMap::new(),
            agent_lines: HashMap::new(),
            prompt_buf: String::new(),
            port,
        }
    }

    /// Build the flat sidebar from current daemon state.
    pub(crate) fn sidebar(&self) -> Vec<SidebarEntry> {
        let mut entries = Vec::new();
        let agent_shell_ids: HashSet<Uuid> =
            self.state.agents.iter().map(|a| a.vex_shell_id).collect();

        // ── Top section: repos with nested agents/shells ───────
        for repo in &self.state.repos {
            entries.push(SidebarEntry::RepoHeader {
                name: repo.name.clone(),
            });

            // Workstreams under this repo
            for ws in &self.state.workstreams {
                if ws.repo == repo.name {
                    entries.push(SidebarEntry::WorkstreamItem {
                        repo: ws.repo.clone(),
                        name: ws.name.clone(),
                    });
                }
            }

            // Agents whose CWD is inside this repo
            for agent in &self.state.agents {
                let cwd = agent.cwd.to_string_lossy();
                let repo_path = repo.path.to_string_lossy();
                if cwd.starts_with(repo_path.as_ref()) {
                    let short = &agent.vex_shell_id.to_string()[..8];
                    let status = if agent.needs_intervention {
                        "NEEDS *"
                    } else {
                        "idle"
                    };
                    entries.push(SidebarEntry::NestedAgent {
                        id: agent.vex_shell_id,
                        label: format!("{} {}", short, status),
                        needs: agent.needs_intervention,
                    });
                }
            }

            // Non-agent shells (can't map to repo without CWD info, skip)
        }

        // ── Divider ────────────────────────────────────────────
        entries.push(SidebarEntry::Divider);

        // ── Agents section ─────────────────────────────────────
        entries.push(SidebarEntry::AgentSectionHeader);
        if self.state.agents.is_empty() {
            entries.push(SidebarEntry::CreateAgentPlaceholder);
        } else {
            for agent in &self.state.agents {
                let short = &agent.vex_shell_id.to_string()[..8];
                let status = if agent.needs_intervention {
                    "NEEDS *"
                } else {
                    "idle"
                };
                entries.push(SidebarEntry::AgentItem {
                    id: agent.vex_shell_id,
                    label: format!("{} {}", short, status),
                    needs: agent.needs_intervention,
                });
            }
        }

        // ── Shells section (non-agent shells only) ─────────────
        let non_agent_shells: Vec<_> = self
            .state
            .shells
            .iter()
            .filter(|s| !agent_shell_ids.contains(&s.id))
            .collect();

        if !non_agent_shells.is_empty() {
            entries.push(SidebarEntry::ShellSectionHeader);
            for shell in non_agent_shells {
                let short = &shell.id.to_string()[..8];
                entries.push(SidebarEntry::ShellItem {
                    id: shell.id,
                    label: format!("{} ({}c)", short, shell.client_count),
                });
            }
        }

        entries
    }

    /// Move cursor down to the next selectable entry.
    pub(crate) fn cursor_down(&mut self) {
        let entries = self.sidebar();
        let mut next = self.cursor + 1;
        while next < entries.len() {
            if entries[next].is_selectable() {
                self.cursor = next;
                return;
            }
            next += 1;
        }
    }

    /// Move cursor up to the previous selectable entry.
    pub(crate) fn cursor_up(&mut self) {
        let entries = self.sidebar();
        if self.cursor == 0 {
            return;
        }
        let mut prev = self.cursor - 1;
        loop {
            if entries[prev].is_selectable() {
                self.cursor = prev;
                return;
            }
            if prev == 0 {
                break;
            }
            prev -= 1;
        }
    }

    /// Clamp cursor after state changes.
    fn clamp_cursor(&mut self) {
        let entries = self.sidebar();
        if entries.is_empty() {
            self.cursor = 0;
            return;
        }
        if self.cursor >= entries.len() {
            self.cursor = entries.len().saturating_sub(1);
        }
        // If on a non-selectable, move to next selectable
        if !entries[self.cursor].is_selectable() {
            let saved = self.cursor;
            self.cursor_down();
            // If cursor_down didn't move (nothing below), try up
            if self.cursor == saved {
                self.cursor_up();
            }
        }
    }

    /// Get the currently selected sidebar entry.
    pub(crate) fn selected(&self) -> Option<SidebarEntry> {
        self.sidebar().get(self.cursor).cloned()
    }

    /// Is the cursor currently in the agents section?
    #[allow(dead_code)]
    pub(crate) fn in_agent_section(&self) -> bool {
        // Walk backwards from cursor to find the nearest section header
        let entries = self.sidebar();
        for i in (0..=self.cursor).rev() {
            match &entries[i] {
                SidebarEntry::AgentSectionHeader => return true,
                SidebarEntry::ShellSectionHeader
                | SidebarEntry::Divider
                | SidebarEntry::RepoHeader { .. } => return false,
                _ => continue,
            }
        }
        false
    }
}

// ── Run loop ───────────────────────────────────────────────────

pub async fn run(port: u16) -> Result<()> {
    let hub = Arc::new(Hub::new(port));
    let mut state_rx = hub.state_rx();
    let command_tx = hub.command_tx();
    let mut event_rx = hub.event_rx();

    let hub_clone = Arc::clone(&hub);
    tokio::spawn(async move {
        if let Err(e) = hub_clone.run().await {
            eprintln!("hub error: {}", e);
        }
    });

    enable_raw_mode()?;
    let mut stdout = io::stdout();
    execute!(stdout, EnterAlternateScreen)?;
    let backend = CrosstermBackend::new(stdout);
    let mut terminal = Terminal::new(backend)?;

    let mut app = App::new(port);

    if let Ok(()) = state_rx.changed().await {
        app.state = state_rx.borrow_and_update().clone();
        app.clamp_cursor();
    }

    let result = run_loop(
        &mut terminal,
        &mut app,
        &mut state_rx,
        &mut event_rx,
        &command_tx,
    )
    .await;

    disable_raw_mode()?;
    execute!(terminal.backend_mut(), LeaveAlternateScreen)?;
    terminal.show_cursor()?;
    result
}

async fn run_loop(
    terminal: &mut Terminal<CrosstermBackend<io::Stdout>>,
    app: &mut App,
    state_rx: &mut watch::Receiver<HubState>,
    event_rx: &mut broadcast::Receiver<FrontendEvent>,
    command_tx: &mpsc::Sender<FrontendCommand>,
) -> Result<()> {
    loop {
        if !app.running {
            break;
        }

        terminal.draw(|f| ui::draw(f, app))?;

        tokio::select! {
            _ = state_rx.changed() => {
                app.state = state_rx.borrow_and_update().clone();
                app.clamp_cursor();
                // Prune stale terminal/conversation data
                let shell_ids: HashSet<Uuid> = app.state.shells.iter().map(|s| s.id).collect();
                let agent_ids: HashSet<Uuid> = app.state.agents.iter().map(|a| a.vex_shell_id).collect();
                app.vt_terminals.retain(|id, _| shell_ids.contains(id));
                app.agent_lines.retain(|id, _| agent_ids.contains(id));
            }
            result = event_rx.recv() => {
                match result {
                    Ok(FrontendEvent::ShellOutput { shell_id, data }) => {
                        if let Some(vt) = app.vt_terminals.get_mut(&shell_id) {
                            vt.process(&data);
                        }
                    }
                    Ok(FrontendEvent::AgentConversationLine { shell_id, line }) => {
                        app.agent_lines.entry(shell_id).or_default().push(line);
                    }
                    Ok(FrontendEvent::AgentWatchEnd { shell_id }) => {
                        app.agent_lines.entry(shell_id).or_default().push("[agent turn complete]".to_string());
                    }
                    Ok(_) => {}
                    Err(broadcast::error::RecvError::Lagged(_)) => continue,
                    Err(_) => break,
                }
            }
            _ = tokio::time::sleep(Duration::from_millis(50)) => {
                while event::poll(Duration::ZERO)? {
                    match event::read()? {
                        Event::Key(key) => handle_key(app, key, command_tx).await?,
                        Event::Resize(cols, rows) => {
                            if let RightPane::ShellTerminal(id) = app.right_pane {
                                if let Some(vt) = app.vt_terminals.get_mut(&id) {
                                    vt.resize(rows, cols);
                                }
                                let _ = command_tx.send(FrontendCommand::ShellResize { id, cols, rows }).await;
                            }
                        }
                        _ => {}
                    }
                }
            }
        }
    }
    Ok(())
}

// ── Key handling ───────────────────────────────────────────────

async fn handle_key(
    app: &mut App,
    key: KeyEvent,
    tx: &mpsc::Sender<FrontendCommand>,
) -> Result<()> {
    if key.code == KeyCode::Char('c') && key.modifiers.contains(KeyModifiers::CONTROL) {
        app.running = false;
        return Ok(());
    }

    // If a dialog is open, route there
    if app.dialog != Dialog::None {
        return handle_dialog_key(app, key, tx).await;
    }

    // If viewing a shell terminal in right pane and it's focused
    if let RightPane::ShellTerminal(id) = app.right_pane {
        // Ctrl+] detaches
        if key.code == KeyCode::Char(']') && key.modifiers.contains(KeyModifiers::CONTROL) {
            let _ = tx.send(FrontendCommand::ShellDetach { id }).await;
            app.right_pane = RightPane::Empty;
            return Ok(());
        }
        // Ctrl+s from agent conversation switches to shell
        // But if we're in shell view, forward all keys except Ctrl+]
        let data = key_to_bytes(key);
        if !data.is_empty() {
            let _ = tx.send(FrontendCommand::ShellInput { id, data }).await;
        }
        return Ok(());
    }

    match key.code {
        KeyCode::Char('q') => app.running = false,

        // Navigation
        KeyCode::Down | KeyCode::Char('j') => app.cursor_down(),
        KeyCode::Up | KeyCode::Char('k') => app.cursor_up(),

        // Enter: open selected item in right pane
        KeyCode::Enter => {
            if let Some(entry) = app.selected() {
                match entry {
                    SidebarEntry::AgentItem { id, .. } | SidebarEntry::NestedAgent { id, .. } => {
                        let _ = tx.send(FrontendCommand::AgentWatch { shell_id: id }).await;
                        app.right_pane = RightPane::AgentConversation(id);
                    }
                    SidebarEntry::ShellItem { id, .. } | SidebarEntry::NestedShell { id, .. } => {
                        open_shell(app, id, tx).await;
                    }
                    SidebarEntry::CreateAgentPlaceholder => {
                        if !app.state.repos.is_empty() {
                            app.dialog = Dialog::SpawnAgent { repo_idx: 0 };
                        } else {
                            app.status_message =
                                Some("register a repo first with `vex repo add`".to_string());
                        }
                    }
                    _ => {}
                }
            }
        }

        // s: open shell for the selected agent, or Ctrl+s from conversation
        KeyCode::Char('s') => {
            if key.modifiers.contains(KeyModifiers::CONTROL) {
                // Ctrl+s from agent conversation pane
                if let RightPane::AgentConversation(id) = app.right_pane {
                    open_shell(app, id, tx).await;
                }
            } else if let Some(
                SidebarEntry::AgentItem { id, .. } | SidebarEntry::NestedAgent { id, .. },
            ) = app.selected()
            {
                open_shell(app, id, tx).await;
            }
        }

        // c: create agent (dialog)
        KeyCode::Char('c') => {
            if app.in_agent_section()
                || matches!(app.selected(), Some(SidebarEntry::CreateAgentPlaceholder))
            {
                if !app.state.repos.is_empty() {
                    app.dialog = Dialog::SpawnAgent { repo_idx: 0 };
                } else {
                    app.status_message = Some("register a repo first".to_string());
                }
            }
        }

        // x: kill selected agent or shell
        KeyCode::Char('x') => {
            if let Some(entry) = app.selected() {
                match entry {
                    SidebarEntry::AgentItem { id, .. } | SidebarEntry::NestedAgent { id, .. } => {
                        let _ = tx.send(FrontendCommand::ShellKill { id }).await;
                        if app.right_pane == RightPane::AgentConversation(id)
                            || app.right_pane == RightPane::ShellTerminal(id)
                        {
                            app.right_pane = RightPane::Empty;
                        }
                        app.status_message =
                            Some(format!("killing agent {}...", &id.to_string()[..8]));
                    }
                    SidebarEntry::ShellItem { id, .. } | SidebarEntry::NestedShell { id, .. } => {
                        let _ = tx.send(FrontendCommand::ShellKill { id }).await;
                        if app.right_pane == RightPane::ShellTerminal(id) {
                            app.right_pane = RightPane::Empty;
                        }
                        app.status_message =
                            Some(format!("killing shell {}...", &id.to_string()[..8]));
                    }
                    _ => {}
                }
            }
        }

        // p: prompt agent (from sidebar or conversation pane)
        KeyCode::Char('p') => {
            let agent_id = match app.selected() {
                Some(SidebarEntry::AgentItem { id, .. } | SidebarEntry::NestedAgent { id, .. }) => {
                    Some(id)
                }
                _ => {
                    if let RightPane::AgentConversation(id) = app.right_pane {
                        Some(id)
                    } else {
                        None
                    }
                }
            };
            if let Some(id) = agent_id {
                app.prompt_buf.clear();
                app.dialog = Dialog::PromptAgent { shell_id: id };
            }
        }

        // +r: new repo dialog (Shift+R)
        KeyCode::Char('R') => {
            app.dialog = Dialog::NewRepo {
                path_buf: String::new(),
                name_buf: String::new(),
                editing_name: false,
            };
        }

        // w: new workstream dialog
        KeyCode::Char('w') => {
            if !app.state.repos.is_empty() {
                app.dialog = Dialog::NewWorkstream {
                    repo_idx: 0,
                    name_buf: String::new(),
                };
            } else {
                app.status_message = Some("register a repo first".to_string());
            }
        }

        // r: refresh
        KeyCode::Char('r') => {
            let _ = tx.send(FrontendCommand::RefreshState).await;
        }

        // ?: help
        KeyCode::Char('?') => {
            app.status_message = Some(
                "Enter:open  s:shell  c:agent  x:kill  p:prompt  R:repo  w:workstream  r:refresh  q:quit"
                    .to_string(),
            );
        }

        KeyCode::Esc => {
            app.status_message = None;
            if app.right_pane != RightPane::Empty {
                app.right_pane = RightPane::Empty;
            }
        }

        _ => {}
    }
    Ok(())
}

async fn open_shell(app: &mut App, id: Uuid, tx: &mpsc::Sender<FrontendCommand>) {
    let (cols, rows) = crossterm::terminal::size().unwrap_or((80, 24));
    // Use roughly half width for the right pane terminal
    let pane_cols = cols / 2;
    app.vt_terminals
        .entry(id)
        .or_insert_with(|| VtTerminal::new(pane_cols, rows.saturating_sub(2)));
    let _ = tx
        .send(FrontendCommand::ShellAttach {
            id,
            cols: pane_cols,
            rows: rows.saturating_sub(2),
        })
        .await;
    app.right_pane = RightPane::ShellTerminal(id);
}

// ── Dialog key handling ────────────────────────────────────────

async fn handle_dialog_key(
    app: &mut App,
    key: KeyEvent,
    tx: &mpsc::Sender<FrontendCommand>,
) -> Result<()> {
    match &mut app.dialog {
        Dialog::SpawnAgent { repo_idx } => match key.code {
            KeyCode::Esc => app.dialog = Dialog::None,
            KeyCode::Down | KeyCode::Char('j') => {
                if *repo_idx + 1 < app.state.repos.len() {
                    *repo_idx += 1;
                }
            }
            KeyCode::Up | KeyCode::Char('k') => {
                *repo_idx = repo_idx.saturating_sub(1);
            }
            KeyCode::Enter => {
                if let Some(repo) = app.state.repos.get(*repo_idx) {
                    let _ = tx
                        .send(FrontendCommand::AgentSpawn {
                            repo: repo.name.clone(),
                            workstream: None,
                        })
                        .await;
                    app.status_message = Some(format!("spawning agent in {}...", repo.name));
                }
                app.dialog = Dialog::None;
            }
            _ => {}
        },

        Dialog::PromptAgent { shell_id } => match key.code {
            KeyCode::Esc => app.dialog = Dialog::None,
            KeyCode::Enter => {
                if !app.prompt_buf.is_empty() {
                    let text = app.prompt_buf.clone();
                    let id = *shell_id;
                    let _ = tx
                        .send(FrontendCommand::AgentPrompt { shell_id: id, text })
                        .await;
                    app.prompt_buf.clear();
                    // Open conversation if not already
                    let _ = tx.send(FrontendCommand::AgentWatch { shell_id: id }).await;
                    app.right_pane = RightPane::AgentConversation(id);
                }
                app.dialog = Dialog::None;
            }
            KeyCode::Backspace => {
                app.prompt_buf.pop();
            }
            KeyCode::Char(c) => app.prompt_buf.push(c),
            _ => {}
        },

        Dialog::NewRepo {
            path_buf,
            name_buf,
            editing_name,
        } => match key.code {
            KeyCode::Esc => app.dialog = Dialog::None,
            KeyCode::Tab => *editing_name = !*editing_name,
            KeyCode::Enter => {
                if !path_buf.is_empty() && !name_buf.is_empty() {
                    let _ = tx
                        .send(FrontendCommand::RepoAdd {
                            name: name_buf.clone(),
                            path: path_buf.clone(),
                        })
                        .await;
                    app.status_message = Some(format!("adding repo '{}'...", name_buf));
                }
                app.dialog = Dialog::None;
            }
            KeyCode::Backspace => {
                if *editing_name {
                    name_buf.pop();
                } else {
                    path_buf.pop();
                }
            }
            KeyCode::Char(c) => {
                if *editing_name {
                    name_buf.push(c);
                } else {
                    path_buf.push(c);
                }
            }
            _ => {}
        },

        Dialog::NewWorkstream { repo_idx, name_buf } => match key.code {
            KeyCode::Esc => app.dialog = Dialog::None,
            KeyCode::Down | KeyCode::Char('j') => {
                if *repo_idx + 1 < app.state.repos.len() {
                    *repo_idx += 1;
                }
            }
            KeyCode::Up | KeyCode::Char('k') => {
                *repo_idx = repo_idx.saturating_sub(1);
            }
            KeyCode::Enter => {
                if !name_buf.is_empty()
                    && let Some(repo) = app.state.repos.get(*repo_idx)
                {
                    let _ = tx
                        .send(FrontendCommand::WorkstreamCreate {
                            repo: repo.name.clone(),
                            name: name_buf.clone(),
                        })
                        .await;
                    app.status_message = Some(format!("creating workstream '{}'...", name_buf));
                }
                app.dialog = Dialog::None;
            }
            KeyCode::Backspace => {
                name_buf.pop();
            }
            KeyCode::Char(c) => name_buf.push(c),
            _ => {}
        },

        Dialog::None => {}
    }
    Ok(())
}

// ── Key to bytes (for shell terminal input) ────────────────────

fn key_to_bytes(key: KeyEvent) -> Vec<u8> {
    match key.code {
        KeyCode::Char(c) => {
            if key.modifiers.contains(KeyModifiers::CONTROL) && c.is_ascii_lowercase() {
                vec![(c as u8) - b'a' + 1]
            } else if key.modifiers.contains(KeyModifiers::CONTROL) {
                vec![]
            } else {
                let mut buf = [0u8; 4];
                c.encode_utf8(&mut buf);
                buf[..c.len_utf8()].to_vec()
            }
        }
        KeyCode::Enter => vec![b'\r'],
        KeyCode::Backspace => vec![0x7f],
        KeyCode::Tab => vec![b'\t'],
        KeyCode::Esc => vec![0x1b],
        KeyCode::Up => b"\x1b[A".to_vec(),
        KeyCode::Down => b"\x1b[B".to_vec(),
        KeyCode::Right => b"\x1b[C".to_vec(),
        KeyCode::Left => b"\x1b[D".to_vec(),
        KeyCode::Home => b"\x1b[H".to_vec(),
        KeyCode::End => b"\x1b[F".to_vec(),
        KeyCode::Delete => b"\x1b[3~".to_vec(),
        _ => vec![],
    }
}