cli/sunbeam/src/code/tui.rs

use std::io;
use std::sync::{Arc, Mutex};

use crossterm::event::{self, Event, KeyCode, KeyEvent, KeyModifiers};
use crossterm::terminal::{self, EnterAlternateScreen, LeaveAlternateScreen};
use crossterm::execute;
use ratatui::backend::CrosstermBackend;
use ratatui::layout::{Constraint, Layout, Rect};
use ratatui::style::{Color, Modifier, Style};
use ratatui::text::{Line, Span, Text};
use ratatui::widgets::{Block, Borders, Paragraph, Wrap};
use ratatui::Terminal;
use tracing_subscriber::fmt::MakeWriter;

// ── Sol color palette ──────────────────────────────────────────────────────

const SOL_YELLOW: Color = Color::Rgb(245, 197, 66);
const SOL_AMBER: Color = Color::Rgb(232, 168, 64);
const SOL_BLUE: Color = Color::Rgb(108, 166, 224);
const SOL_RED: Color = Color::Rgb(224, 88, 88);
const SOL_DIM: Color = Color::Rgb(138, 122, 90);
const SOL_GRAY: Color = Color::Rgb(112, 112, 112);
const SOL_FAINT: Color = Color::Rgb(80, 80, 80);
const SOL_STATUS: Color = Color::Rgb(106, 96, 80);
const SOL_APPROVAL_BG: Color = Color::Rgb(50, 42, 20);
const SOL_APPROVAL_CMD: Color = Color::Rgb(200, 180, 120);

// ── In-memory log buffer for tracing ─────────────────────────────────────

const LOG_BUFFER_CAPACITY: usize = 500;

#[derive(Clone)]
pub struct LogBuffer(Arc<Mutex<Vec<String>>>);

impl LogBuffer {
    pub fn new() -> Self {
        Self(Arc::new(Mutex::new(Vec::new())))
    }

    pub fn lines(&self) -> Vec<String> {
        self.0.lock().unwrap().clone()
    }
}

/// Writer that appends each line to the ring buffer.
pub struct LogBufferWriter(Arc<Mutex<Vec<String>>>);

impl io::Write for LogBufferWriter {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let s = String::from_utf8_lossy(buf);
        let mut lines = self.0.lock().unwrap();
        for line in s.lines() {
            if !line.is_empty() {
                lines.push(line.to_string());
                if lines.len() > LOG_BUFFER_CAPACITY {
                    lines.remove(0);
                }
            }
        }
        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl<'a> MakeWriter<'a> for LogBuffer {
    type Writer = LogBufferWriter;

    fn make_writer(&'a self) -> Self::Writer {
        LogBufferWriter(self.0.clone())
    }
}

// ── Virtual viewport ─────────────────────────────────────────────────────

/// Cached pre-wrapped visual lines for the conversation log.
/// Text is wrapped using `textwrap` when content or width changes.
/// Drawing just slices the visible window — O(viewport), zero wrapping by ratatui.
pub struct Viewport {
    /// Pre-wrapped visual lines (one Line per screen row). Already wrapped to width.
    visual_lines: Vec<Line<'static>>,
    /// Width used for the last wrap pass.
    last_width: u16,
    /// True when log content changed.
    dirty: bool,
}

impl Viewport {
    pub fn new() -> Self {
        Self {
            visual_lines: Vec::new(),
            last_width: 0,
            dirty: true,
        }
    }

    pub fn invalidate(&mut self) {
        self.dirty = true;
    }

    /// Total visual (screen) lines.
    pub fn len(&self) -> u16 {
        self.visual_lines.len() as u16
    }

    /// Rebuild pre-wrapped lines from log entries for a given width.
    pub fn rebuild(&mut self, log: &[LogEntry], width: u16) {
        let w = width.max(1) as usize;
        self.visual_lines.clear();

        for entry in log {
            match entry {
                LogEntry::UserInput(text) => {
                    self.visual_lines.push(Line::from(""));
                    // Wrap user input with "> " prefix
                    let prefixed = format!("> {text}");
                    for wrapped in wrap_styled(&prefixed, w, SOL_DIM, Color::White, 2) {
                        self.visual_lines.push(wrapped);
                    }
                    self.visual_lines.push(Line::from(""));
                }
                LogEntry::AssistantText(text) => {
                    let style = Style::default().fg(SOL_YELLOW);
                    for logical_line in text.lines() {
                        if logical_line.is_empty() {
                            self.visual_lines.push(Line::from(""));
                        } else {
                            for wrapped in textwrap::wrap(logical_line, w) {
                                self.visual_lines.push(Line::styled(wrapped.into_owned(), style));
                            }
                        }
                    }
                }
                LogEntry::ToolSuccess { name, detail } => {
                    self.visual_lines.push(Line::from(vec![
                        Span::styled("  ✓ ", Style::default().fg(SOL_BLUE)),
                        Span::styled(name.clone(), Style::default().fg(SOL_AMBER)),
                        Span::styled(format!(" {detail}"), Style::default().fg(SOL_DIM)),
                    ]));
                }
                LogEntry::ToolExecuting { name, detail } => {
                    self.visual_lines.push(Line::from(vec![
                        Span::styled("  ● ", Style::default().fg(SOL_AMBER)),
                        Span::styled(name.clone(), Style::default().fg(SOL_AMBER)),
                        Span::styled(format!(" {detail}"), Style::default().fg(SOL_DIM)),
                    ]));
                }
                LogEntry::ToolFailed { name, detail } => {
                    self.visual_lines.push(Line::from(vec![
                        Span::styled("  ✗ ", Style::default().fg(SOL_RED)),
                        Span::styled(name.clone(), Style::default().fg(SOL_RED)),
                        Span::styled(format!(" {detail}"), Style::default().fg(SOL_DIM)),
                    ]));
                }
                LogEntry::ToolOutput { lines: output_lines, collapsed } => {
                    let show = if *collapsed { 5 } else { output_lines.len() };
                    let style = Style::default().fg(SOL_GRAY);
                    for line in output_lines.iter().take(show) {
                        self.visual_lines.push(Line::styled(format!("    {line}"), style));
                    }
                    if *collapsed && output_lines.len() > 5 {
                        self.visual_lines.push(Line::styled(
                            format!("    … +{} lines", output_lines.len() - 5),
                            Style::default().fg(SOL_FAINT),
                        ));
                    }
                }
                LogEntry::Status(msg) => {
                    self.visual_lines.push(Line::styled(
                        format!("  [{msg}]"),
                        Style::default().fg(SOL_DIM),
                    ));
                }
                LogEntry::Error(msg) => {
                    let style = Style::default().fg(SOL_RED);
                    for wrapped in textwrap::wrap(&format!("  error: {msg}"), w) {
                        self.visual_lines.push(Line::styled(wrapped.into_owned(), style));
                    }
                }
            }
        }

        self.dirty = false;
        self.last_width = width;
    }

    /// Ensure lines are built for the given width. Rebuilds if width changed.
    pub fn ensure(&mut self, log: &[LogEntry], width: u16) {
        if self.dirty || self.last_width != width {
            self.rebuild(log, width);
        }
    }

    /// Get the visible slice of pre-wrapped lines for the scroll position.
    /// Returns owned lines ready to render — NO wrapping by ratatui.
    pub fn window(&self, height: u16, scroll_offset: u16) -> Vec<Line<'static>> {
        let total = self.visual_lines.len() as u16;
        let max_scroll = total.saturating_sub(height);
        let scroll = if scroll_offset == u16::MAX {
            max_scroll
        } else {
            scroll_offset.min(max_scroll)
        };

        let start = scroll as usize;
        let end = (start + height as usize).min(self.visual_lines.len());
        self.visual_lines[start..end].to_vec()
    }

    pub fn max_scroll(&self, height: u16) -> u16 {
        (self.visual_lines.len() as u16).saturating_sub(height)
    }
}

/// Wrap a "> text" line preserving the dim prefix style on the first line
/// and white text style for content. Returns pre-wrapped visual lines.
fn wrap_styled(text: &str, width: usize, prefix_color: Color, text_color: Color, prefix_len: usize) -> Vec<Line<'static>> {
    let wrapped = textwrap::wrap(text, width);
    let mut lines = Vec::with_capacity(wrapped.len());
    for (i, w) in wrapped.iter().enumerate() {
        let s = w.to_string();
        if i == 0 && s.len() >= prefix_len {
            // First line: split into styled prefix + text
            lines.push(Line::from(vec![
                Span::styled(s[..prefix_len].to_string(), Style::default().fg(prefix_color)),
                Span::styled(s[prefix_len..].to_string(), Style::default().fg(text_color)),
            ]));
        } else {
            lines.push(Line::styled(s, Style::default().fg(text_color)));
        }
    }
    lines
}

// ── Message types for the conversation log ─────────────────────────────────

#[derive(Clone)]
pub enum LogEntry {
    UserInput(String),
    AssistantText(String),
    ToolSuccess { name: String, detail: String },
    ToolExecuting { name: String, detail: String },
    ToolFailed { name: String, detail: String },
    ToolOutput { lines: Vec<String>, collapsed: bool },
    Status(String),
    Error(String),
}

// ── Approval state ─────────────────────────────────────────────────────────

pub struct ApprovalPrompt {
    pub tool_name: String,
    pub command: String,
    pub options: Vec<String>,
    pub selected: usize,
}

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

pub struct App {
    pub log: Vec<LogEntry>,
    pub viewport: Viewport,
    pub input: String,
    pub cursor_pos: usize,
    pub scroll_offset: u16,
    pub project_name: String,
    pub branch: String,
    pub model: String,
    pub input_tokens: u32,
    pub output_tokens: u32,
    pub approval: Option<ApprovalPrompt>,
    pub is_thinking: bool,
    pub sol_status: String,
    pub should_quit: bool,
    pub show_logs: bool,
    pub log_buffer: LogBuffer,
    pub log_scroll: u16,
    pub command_history: Vec<String>,
    pub history_index: Option<usize>,
    pub input_saved: String,
    pub needs_redraw: bool,
}

impl App {
    pub fn new(project_name: &str, branch: &str, model: &str, log_buffer: LogBuffer) -> Self {
        Self {
            log: Vec::new(),
            viewport: Viewport::new(),
            input: String::new(),
            cursor_pos: 0,
            scroll_offset: 0,
            project_name: project_name.into(),
            branch: branch.into(),
            model: model.into(),
            input_tokens: 0,
            output_tokens: 0,
            approval: None,
            is_thinking: false,
            sol_status: String::new(),
            should_quit: false,
            show_logs: false,
            log_buffer,
            log_scroll: u16::MAX,
            command_history: Vec::new(),
            history_index: None,
            input_saved: String::new(),
            needs_redraw: true,
        }
    }

    pub fn push_log(&mut self, entry: LogEntry) {
        self.log.push(entry);
        self.viewport.invalidate();
        self.scroll_offset = u16::MAX;
        self.needs_redraw = true;
    }

    /// Batch-add log entries without per-entry viewport rebuilds.
    pub fn push_logs(&mut self, entries: Vec<LogEntry>) {
        self.log.extend(entries);
        self.viewport.invalidate();
        self.scroll_offset = u16::MAX;
        self.needs_redraw = true;
    }

    /// Resolve the u16::MAX auto-scroll sentinel to the actual max scroll
    /// position. Call before applying relative scroll deltas.
    /// Resolve scroll sentinel AND clamp to valid range. Call before
    /// applying any relative scroll delta.
    pub fn resolve_scroll(&mut self, width: u16, height: u16) {
        self.viewport.ensure(&self.log, width);
        let max = self.viewport.max_scroll(height);
        if self.scroll_offset == u16::MAX {
            self.scroll_offset = max;
        } else {
            self.scroll_offset = self.scroll_offset.min(max);
        }
    }

    /// Load command history from a project's .sunbeam/history file.
    pub fn load_history(&mut self, project_path: &str) {
        let path = std::path::Path::new(project_path).join(".sunbeam").join("history");
        if let Ok(contents) = std::fs::read_to_string(&path) {
            self.command_history = contents.lines().map(String::from).collect();
        }
    }

    /// Save command history to a project's .sunbeam/history file.
    pub fn save_history(&self, project_path: &str) {
        let dir = std::path::Path::new(project_path).join(".sunbeam");
        let _ = std::fs::create_dir_all(&dir);
        let path = dir.join("history");
        // Keep last 500 entries
        let start = self.command_history.len().saturating_sub(500);
        let contents = self.command_history[start..].join("\n");
        let _ = std::fs::write(&path, contents);
    }
}

// ── Rendering ──────────────────────────────────────────────────────────────

pub fn draw(frame: &mut ratatui::Frame, app: &mut App) {
    let area = frame.area();

    // Layout: title (1) + log (flex) + input (3) + status (1)
    let chunks = Layout::vertical([
        Constraint::Length(1),  // title bar
        Constraint::Min(5),    // conversation log
        Constraint::Length(3), // input area
        Constraint::Length(1), // status bar
    ])
    .split(area);

    draw_title_bar(frame, chunks[0], app);

    if app.show_logs {
        draw_debug_log(frame, chunks[1], app);
    } else {
        draw_log(frame, chunks[1], app);
    }

    if let Some(ref approval) = app.approval {
        draw_approval(frame, chunks[2], approval);
    } else {
        draw_input(frame, chunks[2], app);
    }

    draw_status_bar(frame, chunks[3], app);
}

fn draw_title_bar(frame: &mut ratatui::Frame, area: Rect, app: &App) {
    let left = vec![
        Span::styled("sunbeam code", Style::default().fg(SOL_YELLOW).add_modifier(Modifier::BOLD)),
        Span::styled(" · ", Style::default().fg(SOL_FAINT)),
        Span::raw(&app.project_name),
        Span::styled(" · ", Style::default().fg(SOL_FAINT)),
        Span::styled(&app.branch, Style::default().fg(SOL_DIM)),
    ];

    // Right side: model name + sol status
    let right_parts = if app.is_thinking {
        let status = if app.sol_status.is_empty() {
            "generating…"
        } else {
            &app.sol_status
        };
        vec![
            Span::styled(status, Style::default().fg(SOL_AMBER).add_modifier(Modifier::ITALIC)),
            Span::styled(" · ", Style::default().fg(SOL_FAINT)),
            Span::styled(&app.model, Style::default().fg(SOL_DIM)),
        ]
    } else {
        vec![Span::styled(&app.model, Style::default().fg(SOL_DIM))]
    };

    let title_line = Line::from(left);
    frame.render_widget(Paragraph::new(title_line), area);

    let right_line = Line::from(right_parts);
    let right_width = right_line.width() as u16 + 1;
    let right_area = Rect {
        x: area.width.saturating_sub(right_width),
        y: area.y,
        width: right_width,
        height: 1,
    };
    frame.render_widget(Paragraph::new(right_line), right_area);
}

fn draw_log(frame: &mut ratatui::Frame, area: Rect, app: &mut App) {
    // Ensure pre-wrapped lines are built for current width
    app.viewport.ensure(&app.log, area.width);

    // Slice only the visible rows — O(viewport), no wrapping by ratatui
    let window = app.viewport.window(area.height, app.scroll_offset);
    frame.render_widget(Paragraph::new(window), area);
}

fn draw_debug_log(frame: &mut ratatui::Frame, area: Rect, app: &App) {
    let log_lines = app.log_buffer.lines();
    let lines: Vec<Line> = std::iter::once(
        Line::from(Span::styled(
            " debug log (Alt+L to close) ",
            Style::default().fg(SOL_AMBER).add_modifier(Modifier::BOLD),
        )),
    )
    .chain(log_lines.iter().map(|l| {
        let color = if l.contains("ERROR") {
            SOL_RED
        } else if l.contains("WARN") {
            SOL_YELLOW
        } else {
            SOL_GRAY
        };
        Line::from(Span::styled(l.as_str(), Style::default().fg(color)))
    }))
    .collect();

    let total = lines.len() as u16;
    let visible = area.height;
    let max_scroll = total.saturating_sub(visible);
    let scroll = if app.log_scroll == u16::MAX {
        max_scroll
    } else {
        app.log_scroll.min(max_scroll)
    };

    let widget = Paragraph::new(Text::from(lines))
        .wrap(Wrap { trim: false })
        .scroll((scroll, 0));

    frame.render_widget(widget, area);
}

fn draw_input(frame: &mut ratatui::Frame, area: Rect, app: &App) {
    let input_block = Block::default()
        .borders(Borders::TOP)
        .border_style(Style::default().fg(SOL_FAINT));

    let input_text = Line::from(vec![
        Span::styled("> ", Style::default().fg(SOL_DIM)),
        Span::raw(&app.input),
    ]);

    let input_widget = Paragraph::new(input_text)
        .block(input_block)
        .wrap(Wrap { trim: false });

    frame.render_widget(input_widget, area);

    if !app.is_thinking {
        // Only show cursor when not waiting for Sol
        let cursor_x = area.x + 2 + app.cursor_pos as u16;
        let cursor_y = area.y + 1;
        frame.set_cursor_position((cursor_x, cursor_y));
    }
}

fn draw_approval(frame: &mut ratatui::Frame, area: Rect, approval: &ApprovalPrompt) {
    let block = Block::default()
        .borders(Borders::TOP)
        .border_style(Style::default().fg(SOL_FAINT));

    let mut lines = vec![
        Line::from(vec![
            Span::styled("  ⚠ ", Style::default().fg(SOL_YELLOW)),
            Span::styled(&approval.tool_name, Style::default().fg(SOL_YELLOW).add_modifier(Modifier::BOLD)),
            Span::styled(format!("  {}", approval.command), Style::default().fg(SOL_APPROVAL_CMD)),
        ]),
    ];

    for (i, opt) in approval.options.iter().enumerate() {
        let prefix = if i == approval.selected { "  › " } else { "    " };
        let style = if i == approval.selected {
            Style::default().fg(SOL_YELLOW)
        } else {
            Style::default().fg(SOL_DIM)
        };
        lines.push(Line::from(Span::styled(format!("{prefix}{opt}"), style)));
    }

    let widget = Paragraph::new(Text::from(lines))
        .block(block)
        .style(Style::default().bg(SOL_APPROVAL_BG));

    frame.render_widget(widget, area);
}

fn draw_status_bar(frame: &mut ratatui::Frame, area: Rect, app: &App) {
    let left = vec![
        Span::styled(
            format!("  ~/…/{}", app.project_name),
            Style::default().fg(SOL_STATUS),
        ),
        Span::styled(
            format!("  {} ±", app.branch),
            Style::default().fg(SOL_STATUS),
        ),
        Span::styled(
            format!("  {}k in · {}k out", app.input_tokens / 1000, app.output_tokens / 1000),
            Style::default().fg(SOL_STATUS),
        ),
    ];
    frame.render_widget(Paragraph::new(Line::from(left)), area);

    let hint = if app.show_logs { "Alt+L close log" } else { "Alt+L debug log" };
    let right = Span::styled(format!("{hint}  "), Style::default().fg(SOL_FAINT));
    let right_area = Rect {
        x: area.width.saturating_sub(right.width() as u16),
        y: area.y,
        width: right.width() as u16,
        height: 1,
    };
    frame.render_widget(Paragraph::new(Line::from(right)), right_area);
}

// ── Terminal setup/teardown ────────────────────────────────────────────────

pub fn setup_terminal() -> io::Result<Terminal<CrosstermBackend<io::Stdout>>> {
    terminal::enable_raw_mode()?;
    let mut stdout = io::stdout();
    execute!(stdout, EnterAlternateScreen, crossterm::event::EnableMouseCapture)?;
    let backend = CrosstermBackend::new(stdout);
    Terminal::new(backend)
}

pub fn restore_terminal(terminal: &mut Terminal<CrosstermBackend<io::Stdout>>) -> io::Result<()> {
    terminal::disable_raw_mode()?;
    execute!(
        terminal.backend_mut(),
        LeaveAlternateScreen,
        crossterm::event::DisableMouseCapture
    )?;
    terminal.show_cursor()?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_app_creation() {
        let app = App::new("sol", "mainline", "devstral-2", LogBuffer::new());
        assert_eq!(app.project_name, "sol");
        assert!(!app.should_quit);
        assert!(app.log.is_empty());
    }

    #[test]
    fn test_push_log_auto_scrolls() {
        let mut app = App::new("sol", "main", "devstral-2", LogBuffer::new());
        app.scroll_offset = 0;
        app.push_log(LogEntry::Status("test".into()));
        assert_eq!(app.scroll_offset, u16::MAX); // auto-scroll to bottom
    }

    #[test]
    fn test_color_constants() {
        assert!(matches!(SOL_YELLOW, Color::Rgb(245, 197, 66)));
        assert!(matches!(SOL_AMBER, Color::Rgb(232, 168, 64)));
        assert!(matches!(SOL_BLUE, Color::Rgb(108, 166, 224)));
        assert!(matches!(SOL_RED, Color::Rgb(224, 88, 88)));
        // No green in the palette
        assert!(!matches!(SOL_YELLOW, Color::Rgb(_, 255, _)));
        assert!(!matches!(SOL_BLUE, Color::Rgb(_, 255, _)));
    }

    #[test]
    fn test_log_entries_all_variants() {
        let mut app = App::new("test", "main", "devstral-2", LogBuffer::new());
        app.push_log(LogEntry::UserInput("hello".into()));
        app.push_log(LogEntry::AssistantText("response".into()));
        app.push_log(LogEntry::ToolSuccess { name: "file_read".into(), detail: "src/main.rs".into() });
        app.push_log(LogEntry::ToolExecuting { name: "bash".into(), detail: "cargo test".into() });
        app.push_log(LogEntry::ToolFailed { name: "grep".into(), detail: "no matches".into() });
        app.push_log(LogEntry::ToolOutput { lines: vec!["line 1".into(), "line 2".into()], collapsed: true });
        app.push_log(LogEntry::Status("thinking".into()));
        app.push_log(LogEntry::Error("connection lost".into()));
        assert_eq!(app.log.len(), 8);
    }

    #[test]
    fn test_tool_output_collapse_threshold() {
        // Collapsed output shows max 5 lines + "... +N lines"
        let lines: Vec<String> = (0..20).map(|i| format!("line {i}")).collect();
        let entry = LogEntry::ToolOutput { lines: lines.clone(), collapsed: true };
        if let LogEntry::ToolOutput { lines, collapsed } = &entry {
            assert!(lines.len() > 5);
            assert!(*collapsed);
        }
    }

    #[test]
    fn test_approval_prompt() {
        let approval = ApprovalPrompt {
            tool_name: "bash".into(),
            command: "cargo test".into(),
            options: vec![
                "Yes".into(),
                "Yes, always allow bash".into(),
                "No".into(),
            ],
            selected: 0,
        };
        assert_eq!(approval.options.len(), 3);
        assert_eq!(approval.selected, 0);
    }

    #[test]
    fn test_approval_navigation() {
        let mut approval = ApprovalPrompt {
            tool_name: "bash".into(),
            command: "rm -rf".into(),
            options: vec!["Yes".into(), "No".into()],
            selected: 0,
        };
        // Navigate down
        approval.selected = (approval.selected + 1).min(approval.options.len() - 1);
        assert_eq!(approval.selected, 1);
        // Navigate down again (clamped)
        approval.selected = (approval.selected + 1).min(approval.options.len() - 1);
        assert_eq!(approval.selected, 1);
        // Navigate up
        approval.selected = approval.selected.saturating_sub(1);
        assert_eq!(approval.selected, 0);
    }

    #[test]
    fn test_thinking_state() {
        let mut app = App::new("sol", "main", "devstral-2", LogBuffer::new());
        assert!(!app.is_thinking);
        app.is_thinking = true;
        assert!(app.is_thinking);
    }

    #[test]
    fn test_input_cursor() {
        let mut app = App::new("sol", "main", "devstral-2", LogBuffer::new());
        app.input = "hello world".into();
        app.cursor_pos = 5;
        assert_eq!(&app.input[..app.cursor_pos], "hello");
    }

    #[test]
    fn test_token_tracking() {
        let mut app = App::new("sol", "main", "devstral-2", LogBuffer::new());
        app.input_tokens = 1200;
        app.output_tokens = 340;
        assert_eq!(app.input_tokens / 1000, 1);
        assert_eq!(app.output_tokens / 1000, 0);
    }
}