consciousness/src/user/context.rs

use ratatui::{
    layout::Rect,
    style::{Color, Style},
    text::Line,
    Frame,
};

use super::{App, ScreenView, screen_legend};
use super::widgets::{SectionTree, SectionView, section_to_view, pane_block, tree_legend};
use crate::agent::context::{AstNode, NodeBody, Ast};

pub(crate) struct ConsciousScreen {
    agent: std::sync::Arc<crate::agent::Agent>,
    tree: SectionTree,
}

impl ConsciousScreen {
    pub fn new(agent: std::sync::Arc<crate::agent::Agent>) -> Self {
        Self { agent, tree: SectionTree::new() }
    }

    fn read_context_views(&self) -> Vec<SectionView> {
        let ctx = match self.agent.context.try_lock() {
            Ok(ctx) => ctx,
            Err(_) => return Vec::new(),
        };

        let mut views: Vec<SectionView> = Vec::new();

        views.push(section_to_view("System", ctx.system()));
        views.push(section_to_view("Identity", ctx.identity()));
        views.push(section_to_view("Journal", ctx.journal()));

        // Memory nodes extracted from conversation
        let mut mem_children: Vec<SectionView> = Vec::new();
        let mut scored = 0usize;
        let mut unscored = 0usize;
        for node in ctx.conversation() {
            if let AstNode::Leaf(leaf) = node {
                if let NodeBody::Memory { key, score, text } = leaf.body() {
                    if score.is_some() { scored += 1; } else { unscored += 1; }
                    mem_children.push(SectionView {
                        name: format!("mem: {}", key),
                        tokens: node.tokens(),
                        content: text.clone(),
                        children: Vec::new(),
                        status: score.map(|s| format!("{:.2}", s)).unwrap_or_default(),
                    });
                }
            }
        }
        if !mem_children.is_empty() {
            let mem_tokens: usize = mem_children.iter().map(|c| c.tokens).sum();
            views.push(SectionView {
                name: format!("Memory nodes ({})", mem_children.len()),
                tokens: mem_tokens,
                content: String::new(),
                children: mem_children,
                status: format!("{} scored, {} unscored", scored, unscored),
            });
        }

        let conv = ctx.conversation();
        let mut conv_children: Vec<SectionView> = Vec::new();
        for node in conv {
            let mut view = SectionView {
                name: node.label(),
                tokens: node.tokens(),
                content: match node {
                    AstNode::Leaf(leaf) => leaf.body().text().to_string(),
                    _ => String::new(),
                },
                children: match node {
                    AstNode::Branch { children, .. } => children.iter()
                        .map(|c| SectionView {
                            name: c.label(), tokens: c.tokens(),
                            content: match c { AstNode::Leaf(l) => l.body().text().to_string(), _ => String::new() },
                            children: Vec::new(), status: String::new(),
                        }).collect(),
                    _ => Vec::new(),
                },
                status: String::new(),
            };
            // Show memory attribution inline as status text
            if let AstNode::Branch { memory_scores: ms, .. } = node {
                if !ms.is_empty() {
                    let mut attrs: Vec<(&str, f64)> = ms.iter()
                        .map(|(k, v)| (k.as_str(), *v))
                        .collect();
                    attrs.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
                    let parts: Vec<String> = attrs.iter()
                        .map(|(k, s)| format!("{}({:.1})", k, s))
                        .collect();
                    view.status = format!("← {}", parts.join(" "));
                }
            }
            conv_children.push(view);
        }
        let conv_tokens: usize = conv_children.iter().map(|c| c.tokens).sum();
        views.push(SectionView {
            name: format!("Conversation ({} entries)", conv_children.len()),
            tokens: conv_tokens,
            content: String::new(),
            children: conv_children,
            status: String::new(),
        });

        views
    }
}

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_views();
                self.tree.handle_nav(key.code, &context_state, area.height);
            }
        }

        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)));

        let context_state = self.read_context_views();
        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(""));

            self.tree.render_sections(&context_state, &mut lines);

            lines.push(Line::raw(format!("  {:53} {:>6} tokens", "────────", "──────")));
            lines.push(Line::raw(format!("  {:53} {:>6} tokens", "Total", total)));
        } else if let Some(ref info) = app.context_info {
            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)));
        let tool_count = app.agent.state.try_lock()
            .map(|st| st.active_tools.len()).unwrap_or(0);
        lines.push(Line::raw(format!("  Active tools: {}", tool_count)));

        let block = pane_block("context")
            .title_top(Line::from(screen_legend()).left_aligned())
            .title_bottom(tree_legend());

        let widget = super::scroll_pane::ScrollPane::new(&lines).block(block);
        frame.render_stateful_widget(widget, area, &mut self.tree.scroll);
    }
}