use anyhow::Result;
use code2prompt_core::session::Code2PromptSession;
use crossterm::{
execute,
terminal::{EnterAlternateScreen, LeaveAlternateScreen, disable_raw_mode, enable_raw_mode},
};
use ratatui::{
crossterm::event::{KeyCode, KeyEvent, KeyModifiers},
prelude::*,
widgets::*,
};
use std::io::{Stdout, stdout};
use tokio::sync::mpsc;
use crate::clipboard::copy_to_clipboard;
use crate::model::{
AnalysisResults, Cmd, FileTreeInputMode, Message, Model, StatisticsView, Tab, TemplateState,
template::{FocusMode, TemplateFocus, VariableCategory},
};
use crate::token_map::generate_token_map_with_limit;
use crate::utils::{save_template_to_custom_dir, save_to_file};
use crate::widgets::{
FileSelectionWidget, OutputWidget, SettingsWidget, StatisticsByExtensionWidget,
StatisticsOverviewWidget, StatisticsTokenMapWidget, TemplateWidget,
};
use crate::utils::build_file_tree_from_session;
pub struct TuiApp {
model: Model,
terminal: Terminal<CrosstermBackend<Stdout>>,
message_tx: mpsc::UnboundedSender<Message>,
message_rx: mpsc::UnboundedReceiver<Message>,
}
impl TuiApp {
pub fn new(session: Code2PromptSession) -> Result<Self> {
let terminal = init_terminal()?;
let (message_tx, message_rx) = mpsc::unbounded_channel();
let model = Model::new(session);
Ok(Self {
model,
terminal,
message_tx,
message_rx,
})
}
pub async fn run(&mut self) -> Result<()> {
self.handle_message(Message::RefreshFileTree)?;
loop {
let mut messages = Vec::new();
while crossterm::event::poll(std::time::Duration::from_millis(0))? {
if let crossterm::event::Event::Key(key) = crossterm::event::read()?
&& key.kind == crossterm::event::KeyEventKind::Press
{
let ratatui_key = self.convert_crossterm_key(key);
if let Some(message) = self.handle_key_event(ratatui_key) {
if let Some(last_message) = messages.last_mut()
&& self.try_coalesce_messages(last_message, &message)
{
continue; }
messages.push(message);
}
}
}
for message in messages {
self.handle_message(message)?;
}
while let Ok(message) = self.message_rx.try_recv() {
self.handle_message(message)?;
}
let model = self.model.clone();
self.terminal.draw(|frame| {
TuiApp::render_with_model(&model, frame);
})?;
if self.model.should_quit {
break;
}
tokio::time::sleep(tokio::time::Duration::from_millis(1)).await;
}
Ok(())
}
fn render_with_model(model: &Model, frame: &mut Frame) {
let area = frame.area();
let main_layout = Layout::default()
.direction(Direction::Vertical)
.constraints([
Constraint::Length(3), Constraint::Min(0), Constraint::Length(3), ])
.split(area);
Self::render_tab_bar_static(model, frame, main_layout[0]);
match model.current_tab {
Tab::FileTree => {
let widget = FileSelectionWidget::new(model);
let mut state = ();
frame.render_stateful_widget(widget, main_layout[1], &mut state);
}
Tab::Settings => {
let widget = SettingsWidget::new(model);
let mut state = ();
frame.render_stateful_widget(widget, main_layout[1], &mut state);
}
Tab::Statistics => match model.statistics.view {
StatisticsView::Overview => {
let widget = StatisticsOverviewWidget::new(model);
frame.render_widget(widget, main_layout[1]);
}
StatisticsView::TokenMap => {
let widget = StatisticsTokenMapWidget::new(model);
let mut state = ();
frame.render_stateful_widget(widget, main_layout[1], &mut state);
}
StatisticsView::Extensions => {
let widget = StatisticsByExtensionWidget::new(model);
let mut state = ();
frame.render_stateful_widget(widget, main_layout[1], &mut state);
}
},
Tab::Template => {
let widget = TemplateWidget::new(model);
let mut state = TemplateState::from_model(model);
frame.render_stateful_widget(widget, main_layout[1], &mut state);
}
Tab::PromptOutput => {
let widget = OutputWidget::new(model);
let mut state = ();
frame.render_stateful_widget(widget, main_layout[1], &mut state);
}
}
Self::render_status_bar_static(model, frame, main_layout[2]);
}
fn handle_key_event(&self, key: KeyEvent) -> Option<Message> {
if self.model.file_tree_input_mode == FileTreeInputMode::Search
&& self.model.current_tab == Tab::FileTree
{
return self.handle_file_tree_keys(key);
}
if self.model.current_tab == Tab::Template && self.model.template.is_in_editing_mode() {
if key.code == KeyCode::Esc {
return Some(Message::SetTemplateFocusMode(FocusMode::Normal));
}
return self.handle_template_keys(key);
}
match key.code {
KeyCode::Char('q') if key.modifiers.contains(KeyModifiers::CONTROL) => {
return Some(Message::Quit);
}
KeyCode::Esc => return Some(Message::Quit),
KeyCode::Char('1') => return Some(Message::SwitchTab(Tab::FileTree)),
KeyCode::Char('2') => return Some(Message::SwitchTab(Tab::Settings)),
KeyCode::Char('3') => return Some(Message::SwitchTab(Tab::Statistics)),
KeyCode::Char('4') => return Some(Message::SwitchTab(Tab::Template)),
KeyCode::Char('5') => return Some(Message::SwitchTab(Tab::PromptOutput)),
KeyCode::Tab if !key.modifiers.contains(KeyModifiers::SHIFT) => {
let next_tab = match self.model.current_tab {
Tab::FileTree => Tab::Settings,
Tab::Settings => Tab::Statistics,
Tab::Statistics => Tab::Template,
Tab::Template => Tab::PromptOutput,
Tab::PromptOutput => Tab::FileTree,
};
return Some(Message::SwitchTab(next_tab));
}
KeyCode::BackTab | KeyCode::Tab if key.modifiers.contains(KeyModifiers::SHIFT) => {
let prev_tab = match self.model.current_tab {
Tab::FileTree => Tab::PromptOutput,
Tab::Settings => Tab::FileTree,
Tab::Statistics => Tab::Settings,
Tab::Template => Tab::Statistics,
Tab::PromptOutput => Tab::Template,
};
return Some(Message::SwitchTab(prev_tab));
}
_ => {}
}
match self.model.current_tab {
Tab::FileTree => self.handle_file_tree_keys(key),
Tab::Settings => self.handle_settings_keys(key),
Tab::Statistics => self.handle_statistics_keys(key),
Tab::Template => self.handle_template_keys(key),
Tab::PromptOutput => self.handle_prompt_output_keys(key),
}
}
fn handle_file_tree_keys(&self, key: KeyEvent) -> Option<Message> {
if self.model.file_tree_input_mode == FileTreeInputMode::Search {
match key.code {
KeyCode::Esc => Some(Message::ExitSearchMode),
KeyCode::Enter => {
Some(Message::ExitSearchMode)
}
KeyCode::Backspace => {
let mut query = self.model.search_query.clone();
query.pop();
Some(Message::UpdateSearchQuery(query))
}
KeyCode::Char(c) => {
let mut query = self.model.search_query.clone();
query.push(c);
Some(Message::UpdateSearchQuery(query))
}
_ => None,
}
} else {
match key.code {
KeyCode::Up => Some(Message::MoveTreeCursor(-1)),
KeyCode::Down => Some(Message::MoveTreeCursor(1)),
KeyCode::PageUp => Some(Message::MoveTreeCursor(-10)),
KeyCode::PageDown => Some(Message::MoveTreeCursor(10)),
KeyCode::Home => Some(Message::MoveTreeCursor(-9999)),
KeyCode::End => Some(Message::MoveTreeCursor(9999)),
KeyCode::Char(' ') => Some(Message::ToggleFileSelection(self.model.tree_cursor)),
KeyCode::Enter => Some(Message::RunAnalysis),
KeyCode::Right => Some(Message::ExpandDirectory(self.model.tree_cursor)),
KeyCode::Left => Some(Message::CollapseDirectory(self.model.tree_cursor)),
KeyCode::Char('/') => Some(Message::EnterSearchMode),
KeyCode::Char('s') | KeyCode::Char('S') => Some(Message::EnterSearchMode),
KeyCode::Char('r') | KeyCode::Char('R') => Some(Message::RefreshFileTree),
_ => None,
}
}
}
fn handle_settings_keys(&self, key: KeyEvent) -> Option<Message> {
match key.code {
KeyCode::Up => Some(Message::MoveSettingsCursor(-1)),
KeyCode::Down => Some(Message::MoveSettingsCursor(1)),
KeyCode::Char(' ') => Some(Message::ToggleSetting(self.model.settings.settings_cursor)),
KeyCode::Left | KeyCode::Right => {
Some(Message::CycleSetting(self.model.settings.settings_cursor))
}
KeyCode::Enter => Some(Message::RunAnalysis),
_ => None,
}
}
fn handle_statistics_keys(&self, key: KeyEvent) -> Option<Message> {
match key.code {
KeyCode::Enter => Some(Message::RunAnalysis),
KeyCode::Left => Some(Message::CycleStatisticsView(-1)), KeyCode::Right => Some(Message::CycleStatisticsView(1)), KeyCode::Up => Some(Message::ScrollStatistics(-1)),
KeyCode::Down => Some(Message::ScrollStatistics(1)),
KeyCode::PageUp => Some(Message::ScrollStatistics(-5)),
KeyCode::PageDown => Some(Message::ScrollStatistics(5)),
KeyCode::Home => Some(Message::ScrollStatistics(-9999)),
KeyCode::End => Some(Message::ScrollStatistics(9999)),
_ => None,
}
}
fn handle_template_keys(&self, key: KeyEvent) -> Option<Message> {
let is_in_editing_mode = self.model.template.is_in_editing_mode();
let current_focus = self.model.template.get_focus();
if key.code == KeyCode::Esc && is_in_editing_mode {
return Some(Message::SetTemplateFocusMode(FocusMode::Normal));
}
if is_in_editing_mode {
match current_focus {
TemplateFocus::Editor => {
return Some(Message::TemplateEditorInput(key));
}
TemplateFocus::Variables => {
if self.model.template.variables.is_editing() {
match key.code {
KeyCode::Char(c) => return Some(Message::VariableInputChar(c)),
KeyCode::Backspace => return Some(Message::VariableInputBackspace),
KeyCode::Enter => return Some(Message::VariableInputEnter),
KeyCode::Esc => return Some(Message::VariableInputCancel),
_ => return None,
}
} else {
match key.code {
KeyCode::Up => return Some(Message::VariableNavigateUp),
KeyCode::Down => return Some(Message::VariableNavigateDown),
KeyCode::Enter | KeyCode::Char(' ') => {
let variables = self.model.template.get_organized_variables();
if let Some(var) =
variables.get(self.model.template.variables.cursor)
&& var.category == VariableCategory::Missing
{
return Some(Message::VariableStartEditing(var.name.clone()));
}
return None;
}
_ => return None,
}
}
}
_ => {}
}
}
match key.code {
KeyCode::Char('e') | KeyCode::Char('E') => {
return Some(Message::SetTemplateFocus(
TemplateFocus::Editor,
FocusMode::EditingTemplate,
));
}
KeyCode::Char('v') | KeyCode::Char('V') => {
return Some(Message::SetTemplateFocus(
TemplateFocus::Variables,
FocusMode::EditingVariable,
));
}
KeyCode::Char('p') | KeyCode::Char('P') => {
return Some(Message::SetTemplateFocus(
TemplateFocus::Picker,
FocusMode::Normal,
));
}
KeyCode::Char('s') | KeyCode::Char('S') => {
let timestamp = chrono::Utc::now().format("%Y%m%d_%H%M%S");
let filename = format!("custom_template_{}", timestamp);
return Some(Message::SaveTemplate(filename));
}
KeyCode::Char('r') | KeyCode::Char('R') => {
return Some(Message::ReloadTemplate);
}
KeyCode::Enter => {
return Some(Message::RunAnalysis);
}
_ => {}
}
if current_focus == TemplateFocus::Picker {
match key.code {
KeyCode::Up => return Some(Message::TemplatePickerMove(-1)),
KeyCode::Down => return Some(Message::TemplatePickerMove(1)),
KeyCode::Enter | KeyCode::Char('l') | KeyCode::Char('L') | KeyCode::Char(' ') => {
return Some(Message::LoadTemplate);
}
KeyCode::Char('r') | KeyCode::Char('R') => {
return Some(Message::RefreshTemplates);
}
_ => {}
}
}
None
}
fn handle_prompt_output_keys(&self, key: KeyEvent) -> Option<Message> {
match key.code {
KeyCode::Up => Some(Message::ScrollOutput(-1)),
KeyCode::Down => Some(Message::ScrollOutput(1)),
KeyCode::PageUp => Some(Message::ScrollOutput(-10)),
KeyCode::PageDown => Some(Message::ScrollOutput(10)),
KeyCode::Home => Some(Message::ScrollOutput(-9999)),
KeyCode::End => Some(Message::ScrollOutput(9999)),
KeyCode::Char('c') | KeyCode::Char('C') => Some(Message::CopyToClipboard),
KeyCode::Char('s') | KeyCode::Char('S') => {
let timestamp = chrono::Utc::now().format("%Y%m%d_%H%M%S");
let filename = format!("prompt_{}.md", timestamp);
Some(Message::SaveToFile(filename))
}
KeyCode::Enter => Some(Message::RunAnalysis),
_ => None,
}
}
fn handle_message(&mut self, message: Message) -> Result<()> {
let (new_model, cmd) = self.model.update(message);
self.model = new_model;
self.execute_cmd(cmd)?;
Ok(())
}
fn execute_cmd(&mut self, cmd: Cmd) -> Result<()> {
match cmd {
Cmd::None => {
}
Cmd::RefreshFileTree => {
match build_file_tree_from_session(&mut self.model.session) {
Ok(tree) => {
self.model.file_tree_nodes = tree;
self.model.status_message =
"File tree loaded with patterns applied and files auto-expanded"
.to_string();
}
Err(e) => {
self.model.status_message = format!("Error loading files: {}", e);
}
}
}
Cmd::RunAnalysis {
template_content,
user_variables,
} => {
let mut session = self.model.session.clone();
let tx = self.message_tx.clone();
tokio::spawn(async move {
session.config.template_str = template_content;
session.config.template_name = "Custom Template".to_string();
session.config.user_variables = user_variables;
match session.generate_prompt() {
Ok(rendered) => {
let token_map_entries = if rendered.token_count > 0 {
if let Some(files) = session.data.files.as_ref() {
generate_token_map_with_limit(
files,
rendered.token_count,
Some(50),
Some(0.5),
)
} else {
Vec::new()
}
} else {
Vec::new()
};
let result = AnalysisResults {
file_count: rendered.files.len(),
token_count: Some(rendered.token_count),
generated_prompt: rendered.prompt,
token_map_entries,
};
let _ = tx.send(Message::AnalysisComplete(result));
}
Err(e) => {
let _ = tx.send(Message::AnalysisError(e.to_string()));
}
}
});
}
Cmd::CopyToClipboard(content) => match copy_to_clipboard(&content) {
Ok(_) => {
self.model.status_message = "Copied to clipboard!".to_string();
}
Err(e) => {
self.model.status_message = format!("Copy failed: {}", e);
}
},
Cmd::SaveToFile { filename, content } => {
match save_to_file(std::path::Path::new(&filename), &content) {
Ok(_) => {
self.model.status_message = format!("Saved to {}", filename);
}
Err(e) => {
self.model.status_message = format!("Save failed: {}", e);
}
}
}
Cmd::SaveTemplate { filename, content } => {
match save_template_to_custom_dir(std::path::Path::new(&filename), &content) {
Ok(_) => {
self.model.status_message = format!("Template saved as {}", filename);
self.model.template.picker.refresh();
}
Err(e) => {
self.model.status_message = format!("Template save failed: {}", e);
}
}
}
}
Ok(())
}
fn render_tab_bar_static(model: &Model, frame: &mut Frame, area: Rect) {
let tabs = vec![
"1. Selection",
"2. Settings",
"3. Statistics",
"4. Template",
"5. Output",
];
let selected = match model.current_tab {
Tab::FileTree => 0,
Tab::Settings => 1,
Tab::Statistics => 2,
Tab::Template => 3,
Tab::PromptOutput => 4,
};
let tabs_widget = Tabs::new(tabs)
.block(
Block::default()
.borders(Borders::ALL)
.title("Code2Prompt TUI"),
)
.select(selected)
.style(Style::default().fg(Color::White))
.highlight_style(
Style::default()
.fg(Color::Yellow)
.add_modifier(Modifier::BOLD),
);
frame.render_widget(tabs_widget, area);
}
fn render_status_bar_static(model: &Model, frame: &mut Frame, area: Rect) {
let status_text = if !model.status_message.is_empty() {
model.status_message.clone()
} else {
"Tab/Shift+Tab: Switch tabs | 1/2/3/4: Direct tab | Enter: Run Analysis | Esc/Ctrl+Q: Quit".to_string()
};
let status_widget = Paragraph::new(status_text)
.block(Block::default().borders(Borders::ALL))
.style(Style::default().fg(Color::Cyan));
frame.render_widget(status_widget, area);
}
fn convert_crossterm_key(&self, key: crossterm::event::KeyEvent) -> KeyEvent {
use ratatui::crossterm::event::{KeyCode, KeyEventKind, KeyEventState, KeyModifiers};
KeyEvent {
code: match key.code {
crossterm::event::KeyCode::Backspace => KeyCode::Backspace,
crossterm::event::KeyCode::Enter => KeyCode::Enter,
crossterm::event::KeyCode::Left => KeyCode::Left,
crossterm::event::KeyCode::Right => KeyCode::Right,
crossterm::event::KeyCode::Up => KeyCode::Up,
crossterm::event::KeyCode::Down => KeyCode::Down,
crossterm::event::KeyCode::Home => KeyCode::Home,
crossterm::event::KeyCode::End => KeyCode::End,
crossterm::event::KeyCode::PageUp => KeyCode::PageUp,
crossterm::event::KeyCode::PageDown => KeyCode::PageDown,
crossterm::event::KeyCode::Tab => KeyCode::Tab,
crossterm::event::KeyCode::BackTab => KeyCode::BackTab,
crossterm::event::KeyCode::Delete => KeyCode::Delete,
crossterm::event::KeyCode::Insert => KeyCode::Insert,
crossterm::event::KeyCode::F(n) => KeyCode::F(n),
crossterm::event::KeyCode::Char(c) => KeyCode::Char(c),
crossterm::event::KeyCode::Null => KeyCode::Null,
crossterm::event::KeyCode::Esc => KeyCode::Esc,
_ => KeyCode::Null, },
modifiers: KeyModifiers::from_bits_truncate(key.modifiers.bits()),
kind: match key.kind {
crossterm::event::KeyEventKind::Press => KeyEventKind::Press,
crossterm::event::KeyEventKind::Repeat => KeyEventKind::Repeat,
crossterm::event::KeyEventKind::Release => KeyEventKind::Release,
},
state: KeyEventState::from_bits_truncate(key.state.bits()),
}
}
fn try_coalesce_messages(&self, last_message: &mut Message, new_message: &Message) -> bool {
match (last_message, new_message) {
(Message::MoveTreeCursor(delta1), Message::MoveTreeCursor(delta2)) => {
*delta1 += delta2;
true
}
(Message::MoveSettingsCursor(delta1), Message::MoveSettingsCursor(delta2)) => {
*delta1 += delta2;
true
}
(Message::ScrollStatistics(delta1), Message::ScrollStatistics(delta2)) => {
*delta1 += delta2;
true
}
(Message::ScrollOutput(delta1), Message::ScrollOutput(delta2)) => {
*delta1 += delta2;
true
}
(Message::TemplatePickerMove(delta1), Message::TemplatePickerMove(delta2)) => {
*delta1 += delta2;
true
}
_ => false, }
}
}
pub async fn run_tui(session: Code2PromptSession) -> Result<()> {
let mut app = TuiApp::new(session)?;
let result = app.run().await;
restore_terminal()?;
result
}
fn init_terminal() -> Result<Terminal<CrosstermBackend<Stdout>>> {
enable_raw_mode()?;
let mut stdout = stdout();
execute!(stdout, EnterAlternateScreen)?;
let backend = CrosstermBackend::new(stdout);
Terminal::new(backend).map_err(Into::into)
}
fn restore_terminal() -> Result<()> {
disable_raw_mode()?;
execute!(stdout(), LeaveAlternateScreen)?;
Ok(())
}