tiny-trae 0.1.0

An AI coding assistant with tool integration
use anyhow::{Context, Result};
use colored::Colorize;
use log::{debug, info, warn};
use rustyline::{DefaultEditor, error::ReadlineError};
use serde_json::Value;
use std::collections::HashMap;
use std::io::{self, Write};
use uuid::Uuid;

use crate::config::Config;
use crate::claude::{ClaudeClient, Message, ContentBlockWithTools, ChatResponseWithTools};
use crate::tools::{ToolSystem, ToolResult};

pub struct ChatInterface {
    claude_client: ClaudeClient,
    tool_system: ToolSystem,
    conversation_history: Vec<Message>,
    config: Config,
    editor: DefaultEditor,
}

impl ChatInterface {
    pub async fn new(config: Config) -> Result<Self> {
        config.validate()?;
        
        let claude_client = ClaudeClient::new(config.clone())?;
        let tool_system = ToolSystem::new(config.clone());
        
        let mut editor = DefaultEditor::new()
            .context("Failed to create readline editor")?;
        
        // Load history if exists
        if editor.load_history("chat_history.txt").is_err() {
            debug!("No existing chat history found");
        }

        Ok(Self {
            claude_client,
            tool_system,
            conversation_history: vec![
                Message {
                    role: "system".to_string(),
                    content: config.system.prompt.clone(),
                }
            ],
            config,
            editor,
        })
    }

    pub async fn run(&mut self) -> Result<()> {
        println!("{}", "🚀 Tiny Trae AI Coding Agent is ready!".green().bold());
        println!("{}", "Type 'help' for available commands, 'exit' to quit.".yellow());
        println!("{}", "".repeat(60).cyan());

        loop {
            let input = match self.editor.readline(&format!("{}", "🤖 > ".cyan())) {
                Ok(line) => {
                    self.editor.add_history_entry(line.as_str());
                    line
                }
                Err(ReadlineError::Interrupted) => {
                    println!("{}", "Interrupted. Type 'exit' to quit.".yellow());
                    continue;
                }
                Err(ReadlineError::Eof) => {
                    println!("{}", "EOF encountered. Exiting...".yellow());
                    break;
                }
                Err(err) => {
                    eprintln!("Error reading input: {}", err);
                    continue;
                }
            };

            let input = input.trim();
            
            if input.is_empty() {
                continue;
            }

            match input {
                "exit" | "quit" => {
                    println!("{}", "Goodbye! 👋".green());
                    break;
                }
                "help" => {
                    self.show_help();
                    continue;
                }
                "clear" => {
                    self.conversation_history.clear();
                    println!("{}", "Conversation history cleared.".green());
                    continue;
                }
                "history" => {
                    self.show_history();
                    continue;
                }
                "status" => {
                    self.show_status();
                    continue;
                }
                _ => {
                    if let Err(e) = self.process_user_input(input).await {
                        eprintln!("{}: {}", "Error".red(), e);
                    }
                }
            }
        }

        // Save history
        if let Err(e) = self.editor.save_history("chat_history.txt") {
            warn!("Failed to save chat history: {}", e);
        }

        Ok(())
    }

    async fn process_user_input(&mut self, input: &str) -> Result<()> {
        // Add user message to history
        self.conversation_history.push(Message {
            role: "user".to_string(),
            content: input.to_string(),
        });

        println!("{}", "🤔 Thinking...".yellow());

        // Get tool definitions
        let tools = self.tool_system.get_tool_definitions();
        
        // Send message to Claude with tools
        let response = self.claude_client
            .chat_with_tools(self.conversation_history.clone(), tools)
            .await?;

        // Process response
        let assistant_message = self.process_claude_response(response).await?;

        // Add assistant response to history
        self.conversation_history.push(Message {
            role: "assistant".to_string(),
            content: assistant_message,
        });

        Ok(())
    }

    async fn process_claude_response(&mut self, response: ChatResponseWithTools) -> Result<String> {
        let mut assistant_message = String::new();
        let mut tool_results = Vec::new();

        for content_block in response.content {
            match content_block {
                ContentBlockWithTools::Text { text } => {
                    if !text.trim().is_empty() {
                        println!("{}", text.green());
                        assistant_message.push_str(&text);
                    }
                }
                ContentBlockWithTools::ToolUse { id, name, input } => {
                    println!("{}: {}", "🔧 Using tool".blue().bold(), name.cyan());
                    
                    let tool_result = self.tool_system.execute_tool(&name, input).await?;
                    
                    if tool_result.success {
                        println!("{}: {}", "✅ Tool succeeded".green(), name.cyan());
                        if !tool_result.output.trim().is_empty() {
                            println!("{}", tool_result.output.dimmed());
                        }
                    } else {
                        println!("{}: {}", "❌ Tool failed".red(), name.cyan());
                        if let Some(error) = &tool_result.error {
                            println!("{}: {}", "Error".red(), error.red());
                        }
                    }

                    tool_results.push(Message {
                        role: "user".to_string(),
                        content: format!("Tool {} result: {}", name, serde_json::to_string(&tool_result)?),
                    });
                }
            }
        }

        // If there were tool calls, send the results back to Claude
        if !tool_results.is_empty() {
            println!("{}", "🔄 Processing tool results...".yellow());
            
            // Add tool results to conversation
            self.conversation_history.extend(tool_results);
            
            // Get Claude's response to the tool results
            let follow_up_response = self.claude_client
                .chat_with_tools(self.conversation_history.clone(), vec![])
                .await?;

            let follow_up_message = Box::pin(self.process_claude_response(follow_up_response)).await?;
            assistant_message.push_str(&follow_up_message);
        }

        Ok(assistant_message)
    }

    fn show_help(&self) {
        println!("{}", "Available commands:".cyan().bold());
        println!("  {}  - Show this help message", "help".green());
        println!("  {}  - Clear conversation history", "clear".green());
        println!("  {}  - Show conversation history", "history".green());
        println!("  {}  - Show system status", "status".green());
        println!("  {}  - Exit the application", "exit".green());
        println!("\n{}", "Available tools:".cyan().bold());
        let tools = self.tool_system.get_tool_definitions();
        for tool in tools {
            println!("  {} - {}", tool.function.name.green(), tool.function.description);
        }
    }

    fn show_history(&self) {
        if self.conversation_history.is_empty() {
            println!("{}", "No conversation history.".yellow());
            return;
        }

        println!("{}", "Conversation History:".cyan().bold());
        println!("{}", "".repeat(60).cyan());
        
        for (i, message) in self.conversation_history.iter().enumerate() {
            let role_color = if message.role == "user" { 
                "👤 User".blue().bold()
            } else { 
                "🤖 Assistant".green().bold()
            };
            
            println!("{} ({}): {}", role_color, i + 1, message.content);
            println!("{}", "".repeat(60).dimmed());
        }
    }

    fn show_status(&self) {
        println!("{}", "System Status:".cyan().bold());
        println!("  Model: {}", self.config.model.version.green());
        println!("  Max Tokens: {}", self.config.model.max_tokens.to_string().green());
        println!("  Temperature: {}", self.config.model.temperature.to_string().green());
        println!("  Workspace: {}", self.config.workspace.root_path.green());
        println!("  Conversation Length: {}", self.conversation_history.len().to_string().green());
        
        let tools = self.tool_system.get_tool_definitions();
        println!("  Available Tools: {}", tools.len().to_string().green());
        for tool in tools {
            println!("    - {}", tool.function.name.cyan());
        }
    }
}