consciousness/src/user/context.rs

// context_screen.rs — F2 context/debug overlay
//
// Full-screen overlay showing model info, context window breakdown,
// and runtime state. Supports tree navigation with expand/collapse.

use ratatui::{
    layout::Rect,
    style::{Color, Modifier, Style},
    text::Line,
    widgets::{Block, Borders, Paragraph, Wrap},
    Frame,
    crossterm::event::KeyCode,
};

use super::{App, ScreenView, screen_legend};
use crate::agent::context::ContextSection;

pub(crate) struct ConsciousScreen {
    agent: std::sync::Arc<tokio::sync::Mutex<crate::agent::Agent>>,
    scroll: u16,
    selected: Option<usize>,
    expanded: std::collections::HashSet<usize>,
}

impl ConsciousScreen {
    pub fn new(agent: std::sync::Arc<tokio::sync::Mutex<crate::agent::Agent>>) -> Self {
        Self { agent, scroll: 0, selected: None, expanded: std::collections::HashSet::new() }
    }

    fn read_context_state(&self) -> Vec<ContextSection> {
        match self.agent.try_lock() {
            Ok(ag) => ag.context_state_summary(None),
            Err(_) => Vec::new(),
        }
    }

    fn item_count(&self, context_state: &[ContextSection]) -> usize {
        fn count_section(section: &ContextSection, expanded: &std::collections::HashSet<usize>, idx: &mut usize) -> usize {
            let my_idx = *idx;
            *idx += 1;
            let mut total = 1;
            if expanded.contains(&my_idx) {
                for child in &section.children {
                    total += count_section(child, expanded, idx);
                }
            }
            total
        }
        let mut idx = 0;
        let mut total = 0;
        for section in context_state {
            total += count_section(section, &self.expanded, &mut idx);
        }
        total
    }

    fn scroll_to_selected(&mut self, _item_count: usize) {
        let header_lines = 8u16;
        if let Some(sel) = self.selected {
            let sel_line = header_lines + sel as u16;
            if sel_line < self.scroll + 2 {
                self.scroll = sel_line.saturating_sub(2);
            } else if sel_line > self.scroll + 30 {
                self.scroll = sel_line.saturating_sub(15);
            }
        }
    }

    fn render_section(
        &self,
        section: &ContextSection,
        depth: usize,
        lines: &mut Vec<Line>,
        idx: &mut usize,
    ) {
        let my_idx = *idx;
        let selected = self.selected == Some(my_idx);
        let expanded = self.expanded.contains(&my_idx);
        let has_children = !section.children.is_empty();
        let has_content = !section.content.is_empty();
        let expandable = has_children || has_content;

        let indent = "  ".repeat(depth + 1);
        let marker = if !expandable { " " } else if expanded { "▼" } else { "▶" };
        let label = format!("{}{} {:30} {:>6} tokens", indent, marker, section.name, section.tokens);
        let style = if selected {
            Style::default().fg(Color::Yellow).add_modifier(Modifier::BOLD)
        } else {
            Style::default()
        };
        lines.push(Line::styled(label, style));
        *idx += 1;

        if expanded {
            if has_children {
                for child in &section.children {
                    self.render_section(child, depth + 1, lines, idx);
                }
            } else if has_content {
                let content_indent = format!("{}    │ ", "  ".repeat(depth + 1));
                let content_lines: Vec<&str> = section.content.lines().collect();
                let show = content_lines.len().min(50);
                for line in &content_lines[..show] {
                    lines.push(Line::styled(
                        format!("{}{}", content_indent, line),
                        Style::default().fg(Color::DarkGray),
                    ));
                }
                if content_lines.len() > 50 {
                    lines.push(Line::styled(
                        format!("{}... ({} more lines)", content_indent, content_lines.len() - 50),
                        Style::default().fg(Color::DarkGray),
                    ));
                }
            }
        }
    }
}

impl ScreenView for ConsciousScreen {
    fn label(&self) -> &'static str { "conscious" }

    fn tick(&mut self, frame: &mut Frame, area: Rect,
            events: &[ratatui::crossterm::event::Event], app: &mut App) {
        for event in events {
            if let ratatui::crossterm::event::Event::Key(key) = event {
                if key.kind != ratatui::crossterm::event::KeyEventKind::Press { continue; }
                let context_state = self.read_context_state();
                let item_count = self.item_count(&context_state);

                match key.code {
                KeyCode::Up => {
                    self.selected = Some(self.selected.unwrap_or(0).saturating_sub(1));
                    self.scroll_to_selected(item_count);
                }
                KeyCode::Down => {
                    let max = item_count.saturating_sub(1);
                    self.selected = Some(self.selected.map_or(0, |s| (s + 1).min(max)));
                    self.scroll_to_selected(item_count);
                }
                KeyCode::Right | KeyCode::Enter => {
                    if let Some(sel) = self.selected {
                        self.expanded.insert(sel);
                    }
                }
                KeyCode::Left => {
                    if let Some(sel) = self.selected {
                        self.expanded.remove(&sel);
                    }
                }
                KeyCode::PageUp => { self.scroll = self.scroll.saturating_sub(20); }
                KeyCode::PageDown => { self.scroll += 20; }
                _ => {}
            }
            }
        }

        // Draw
        let mut lines: Vec<Line> = Vec::new();
        let section_style = Style::default().fg(Color::Yellow);

        lines.push(Line::styled("── Model ──", section_style));
        let model_display = app.context_info.as_ref()
            .map_or_else(|| app.status.model.clone(), |i| i.model.clone());
        lines.push(Line::raw(format!("  Current: {}", model_display)));
        if let Some(ref info) = app.context_info {
            lines.push(Line::raw(format!("  Backend: {}", info.backend)));
            lines.push(Line::raw(format!("  Prompt:  {}", info.prompt_file)));
            lines.push(Line::raw(format!("  Available: {}", info.available_models.join(", "))));
        }
        lines.push(Line::raw(""));

        lines.push(Line::styled("── Context State ──", section_style));
        lines.push(Line::raw(format!("  Prompt tokens: {}K", app.status.prompt_tokens / 1000)));
        if !app.status.context_budget.is_empty() {
            lines.push(Line::raw(format!("  Budget: {}", app.status.context_budget)));
        }
        let context_state = self.read_context_state();
        if !context_state.is_empty() {
            let total: usize = context_state.iter().map(|s| s.tokens).sum();
            lines.push(Line::raw(""));
            lines.push(Line::styled(
                "  (↑/↓ select, →/Enter expand, ← collapse, PgUp/PgDn scroll)",
                Style::default().fg(Color::DarkGray),
            ));
            lines.push(Line::raw(""));

            let mut flat_idx = 0usize;
            for section in &context_state {
                self.render_section(section, 0, &mut lines, &mut flat_idx);
            }

            lines.push(Line::raw(format!("  {:23} {:>6} tokens", "────────", "──────")));
            lines.push(Line::raw(format!("  {:23} {:>6} tokens", "Total", total)));
        } else if let Some(ref info) = app.context_info {
            lines.push(Line::raw(format!("  System prompt:   {:>6} chars", info.system_prompt_chars)));
            lines.push(Line::raw(format!("  Context message: {:>6} chars", info.context_message_chars)));
        }
        lines.push(Line::raw(""));

        lines.push(Line::styled("── Runtime ──", section_style));
        lines.push(Line::raw(format!(
            "  DMN: {} ({}/{})",
            app.status.dmn_state, app.status.dmn_turns, app.status.dmn_max_turns,
        )));
        lines.push(Line::raw(format!("  Reasoning: {}", app.reasoning_effort)));
        lines.push(Line::raw(format!("  Running processes: {}", app.running_processes)));
        lines.push(Line::raw(format!("  Active tools: {}", app.active_tools.lock().unwrap().len())));

        let block = Block::default()
            .title_top(Line::from(screen_legend()).left_aligned())
            .title_top(Line::from(" context ").right_aligned())
            .borders(Borders::ALL)
            .border_style(Style::default().fg(Color::Cyan));

        let para = Paragraph::new(lines)
            .block(block)
            .wrap(Wrap { trim: false })
            .scroll((self.scroll, 0));

        frame.render_widget(para, area);
    }
}