"""
Enhanced Output Formatter

This module provides advanced formatting utilities for displaying agent responses
in a beautiful, interactive format with colors, animations, and user-friendly features.
"""

import json
import time
import sys
import threading
from typing import Dict, Any, List, Optional
from datetime import datetime
import os


class EnhancedOutputFormatter:
    """Enhanced formatter with advanced visual features and user experience"""
    
    def __init__(self, minimal: bool = False, typing_delay: float = 0.01, enable_animations: bool = True):
        # Enhanced color palette with more vibrant colors
        self.colors = {
            'header': '\033[95m',      # Magenta
            'blue': '\033[94m',        # Blue
            'cyan': '\033[96m',        # Cyan
            'green': '\033[92m',       # Green
            'yellow': '\033[93m',      # Yellow
            'red': '\033[91m',         # Red
            'purple': '\033[35m',      # Purple
            'orange': '\033[33m',      # Orange
            'pink': '\033[95m',        # Pink
            'bold': '\033[1m',         # Bold
            'underline': '\033[4m',    # Underline
            'grey': '\033[90m',        # Grey
            'light_grey': '\033[37m',  # Light Grey
            'end': '\033[0m',          # Reset
            # Background colors
            'bg_blue': '\033[44m',
            'bg_green': '\033[42m',
            'bg_yellow': '\033[43m',
            'bg_red': '\033[41m',
            'bg_cyan': '\033[46m',
        }
        
        self.minimal = minimal
        self.typing_delay = typing_delay
        self.enable_animations = enable_animations
        self._section_lines = 0
        self._progress_active = False
        self._progress_thread = None
        
        # Check if terminal supports colors
        self.supports_colors = self._check_color_support()
        
        # Animation frames for loading indicators
        self.spinner_frames = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏']
        self.dot_frames = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏']
        self.bar_frames = ['▰▱▱▱▱▱▱▱', '▰▰▱▱▱▱▱▱', '▰▰▰▱▱▱▱▱', '▰▰▰▰▱▱▱▱', 
                          '▰▰▰▰▰▱▱▱', '▰▰▰▰▰▰▱▱', '▰▰▰▰▰▰▰▱', '▰▰▰▰▰▰▰▰']

    def _check_color_support(self) -> bool:
        """Check if terminal supports colors"""
        if not hasattr(sys.stdout, 'isatty'):
            return False
        if not sys.stdout.isatty():
            return False
        if 'NO_COLOR' in os.environ:
            return False
        return True

    def _get_color(self, color_name: str) -> str:
        """Get color code, return empty string if colors not supported"""
        if not self.supports_colors:
            return ""
        return self.colors.get(color_name, "")

    def clear_screen(self):
        """Clear the terminal screen"""
        os.system('cls' if os.name == 'nt' else 'clear')

    def print_banner(self, title: str = "AI Agent System"):
        """Print a beautiful banner"""
        banner = f"""
{self._get_color('cyan')}{'='*70}{self._get_color('end')}
{self._get_color('bold')}{self._get_color('purple')}{'🤖 ' + title + ' 🤖':^70}{self._get_color('end')}
{self._get_color('cyan')}{'='*70}{self._get_color('end')}
"""
        print(banner)

    def start_progress(self, message: str = "Processing..."):
        """Start an animated progress indicator"""
        if not self.enable_animations:
            print(f"{self._get_color('blue')}🔄 {message}{self._get_color('end')}")
            return
            
        self._progress_active = True
        self._progress_thread = threading.Thread(target=self._animate_progress, args=(message,))
        self._progress_thread.daemon = True
        self._progress_thread.start()

    def stop_progress(self):
        """Stop the progress animation"""
        if self._progress_active:
            self._progress_active = False
            if self._progress_thread:
                self._progress_thread.join(timeout=0.1)
            # Clear the progress line
            sys.stdout.write('\r' + ' ' * 80 + '\r')
            sys.stdout.flush()

    def _animate_progress(self, message: str):
        """Animate progress indicator"""
        frame_idx = 0
        while self._progress_active:
            frame = self.spinner_frames[frame_idx % len(self.spinner_frames)]
            progress_bar = self.bar_frames[frame_idx % len(self.bar_frames)]
            
            sys.stdout.write(f'\r{self._get_color("cyan")}{frame} {message} {progress_bar}{self._get_color("end")}')
            sys.stdout.flush()
            
            time.sleep(0.1)
            frame_idx += 1

    def typewriter_print(self, text: str, delay: float = None, color: str = None):
        """Enhanced typewriter print with color support"""
        if delay is None:
            delay = self.typing_delay
            
        color_code = self._get_color(color) if color else ""
        end_code = self._get_color('end') if color_code else ""
        
        for char in text:
            sys.stdout.write(f"{color_code}{char}{end_code}")
            sys.stdout.flush()
            if char != '\n':
                time.sleep(delay)
            else:
                self._section_lines += 1
        if text and text[-1] != '\n':
            sys.stdout.write('\n')
            sys.stdout.flush()
            self._section_lines += 1

    def section_start(self, title: str, color: str = 'blue', icon: str = None):
        """Enhanced section header with icons and better styling"""
        self._section_lines = 0
        color_code = self._get_color(color)
        bold_code = self._get_color('bold')
        end_code = self._get_color('end')
        
        if icon:
            header = f"{bold_code}{color_code}{icon} {title}{end_code}"
        else:
            header = f"{bold_code}{color_code}{title}{end_code}"
            
        # Add subtle border
        border = f"{self._get_color('light_grey')}{'─' * (len(title) + 4)}{self._get_color('end')}"
        
        sys.stdout.write(f"{header}\n{border}\n")
        sys.stdout.flush()
        self._section_lines += 2

    def section_end(self, keep_title: bool = True):
        """Clear section content, keeping only the title line if requested"""
        for _ in range(self._section_lines - (2 if keep_title else 0)):
            sys.stdout.write('\033[F')  # Move cursor up
            sys.stdout.write('\033[K')  # Clear line
        sys.stdout.flush()
        self._section_lines = 2 if keep_title else 0

    def format_agent_response(self, response) -> str:
        """Enhanced agent response formatting with better visual hierarchy"""
        output = []
        
        # Enhanced header with gradient effect
        output.append(self._format_enhanced_header("🤖 AI Agent Response"))
        output.append("")
        
        # Execution info with progress indicators
        output.append(self._format_section_header("📊 Execution Summary", "cyan"))
        output.append(self._format_execution_summary(response))
        output.append("")
        
        # Reasoning steps with better visual flow
        if response.reasoning_steps:
            output.append(self._format_section_header("🧠 Reasoning Process", "purple"))
            for i, step in enumerate(response.reasoning_steps, 1):
                output.append(self._format_enhanced_reasoning_step(i, step))
            output.append("")
        
        # Tool output with enhanced formatting
        if response.tool_output:
            output.append(self._format_section_header("🔧 Tool Output", "yellow"))
            output.append(self._format_enhanced_tool_output(response.tool_output))
            output.append("")
        
        # Final response with better typography
        output.append(self._format_section_header("💬 Final Response", "green"))
        output.append(self._format_enhanced_text_content(response.final_response))
        output.append("")

        # Final Result section with highlight
        final_result = self._extract_final_result(response.tool_output)
        if final_result is not None:
            output.append(self._format_section_header("🎯 Final Result", "orange"))
            output.append(self._format_final_result(final_result))
            output.append("")
        
        # Memory updates with better organization
        if hasattr(response, 'memory_updates') and response.memory_updates:
            output.append(self._format_section_header("💾 Memory Updates", "blue"))
            for update in response.memory_updates:
                output.append(f"  {self._get_color('green')}•{self._get_color('end')} {update}")
            output.append("")
        
        output.append(self._format_enhanced_footer())
        
        return "\n".join(output)

    def _format_header(self, title: str) -> str:
        """Format main header (backward compatibility)"""
        return self._format_enhanced_header(title)

    def _format_enhanced_header(self, title: str) -> str:
        """Enhanced header with gradient-like effect"""
        border_top = f"{self._get_color('cyan')}{'╭' + '─' * 58 + '╮'}{self._get_color('end')}"
        title_line = f"{self._get_color('bold')}{self._get_color('purple')}{'│' + title.center(58) + '│'}{self._get_color('end')}"
        border_bottom = f"{self._get_color('cyan')}{'╰' + '─' * 58 + '╯'}{self._get_color('end')}"
        return f"{border_top}\n{title_line}\n{border_bottom}"

    def _format_section_header(self, title: str, color: str = 'blue') -> str:
        """Enhanced section header"""
        color_code = self._get_color(color)
        bold_code = self._get_color('bold')
        end_code = self._get_color('end')
        return f"{bold_code}{color_code}{title}{end_code}"

    def _format_execution_summary(self, response) -> str:
        """Enhanced execution summary with visual indicators"""
        summary = []
        
        # Time with color coding based on duration
        time_color = 'green' if response.execution_time < 5 else ('yellow' if response.execution_time < 15 else 'red')
        summary.append(f"  {self._get_color('blue')}⏱️  Execution Time:{self._get_color('end')} {self._get_color(time_color)}{response.execution_time:.2f} seconds{self._get_color('end')}")
        
        # Tool used with status indicator
        tool_status = "✅" if response.tool_used else "❌"
        summary.append(f"  {self._get_color('blue')}🔧 Tool Used:{self._get_color('end')} {tool_status} {response.tool_used or 'None'}")
        
        # Reasoning steps with visual count
        step_count = len(response.reasoning_steps)
        step_color = 'green' if step_count <= 5 else ('yellow' if step_count <= 10 else 'red')
        summary.append(f"  {self._get_color('blue')}🧠 Reasoning Steps:{self._get_color('end')} {self._get_color(step_color)}{step_count}{self._get_color('end')}")
        
        return "\n".join(summary)

    def _format_reasoning_step(self, step_num: int, step: Dict[str, Any]) -> str:
        """Format a reasoning step (backward compatibility)"""
        return self._format_enhanced_reasoning_step(step_num, step)

    def _format_enhanced_reasoning_step(self, step_num: int, step: Dict[str, Any]) -> str:
        """Enhanced reasoning step with better visual hierarchy"""
        step_type = step.get('step_type', 'unknown')
        if hasattr(step_type, "name"):
            step_type_str = step_type.name
        else:
            step_type_str = str(step_type)
        description = step.get('description', 'No description')
        timestamp = step.get('timestamp', time.time())
        
        # Format timestamp
        dt = datetime.fromtimestamp(timestamp)
        time_str = dt.strftime("%H:%M:%S.%f")[:-3]
        
        # Enhanced color coding by step type
        step_colors = {
            'analysis': 'cyan',
            'tool_selection': 'yellow',
            'execution': 'green',
            'synthesis': 'purple',
            'error': 'red'
        }
        
        color = step_colors.get(step_type_str.lower(), 'blue')
        color_code = self._get_color(color)
        bold_code = self._get_color('bold')
        end_code = self._get_color('end')
        
        # Step number with visual indicator
        step_indicator = f"{bold_code}{color_code}●{end_code}"
        
        output = f"  {step_indicator} {bold_code}{color_code}Step {step_num}: {step_type_str.upper()}{end_code}"
        output += f"\n    {self._get_color('grey')}📝{self._get_color('end')} {description}"
        output += f"\n    {self._get_color('grey')}🕐{self._get_color('end')} {time_str}"
        
        # Add details if present
        details = step.get('details', {})
        if details:
            output += f"\n    {self._get_color('grey')}📋{self._get_color('end')} Details: {json.dumps(details, indent=6)}"
        
        return output

    def _format_tool_output(self, tool_output: Any) -> str:
        """Format tool output (backward compatibility)"""
        return self._format_enhanced_tool_output(tool_output)

    def _format_enhanced_tool_output(self, tool_output: Any) -> str:
        """Enhanced tool output formatting"""
        if isinstance(tool_output, dict):
            if 'success' in tool_output:
                success_indicator = f"{self._get_color('green')}✅{self._get_color('end')}" if tool_output['success'] else f"{self._get_color('red')}❌{self._get_color('end')}"
                status_color = 'green' if tool_output['success'] else 'red'
                
                output = f"  {success_indicator} {self._get_color('bold')}{self._get_color(status_color)}Status:{self._get_color('end')} {'Success' if tool_output['success'] else 'Failed'}"
                
                if 'message' in tool_output:
                    output += f"\n  {self._get_color('blue')}📄 Message:{self._get_color('end')} {tool_output['message']}"
                if 'error' in tool_output:
                    output += f"\n  {self._get_color('red')}🚨 Error:{self._get_color('end')} {tool_output['error']}"
                if 'diff' in tool_output and tool_output.get('comparison_type') in ('unified_diff', 'comprehensive'):
                    output += f"\n  {self._get_color('yellow')}📝 Diff:{self._get_color('end')}\n{self._colorize_diff(tool_output['diff'])}"
                if 'content' in tool_output:
                    output += f"\n  {self._get_color('cyan')}📝 Content:{self._get_color('end')}\n{self._indent_text(str(tool_output['content'])[:500])}..."
                if 'extracted_content' in tool_output:
                    output += f"\n  {self._get_color('purple')}🔍 Extracted Content:{self._get_color('end')}\n{self._format_extracted_content(tool_output['extracted_content'])}"
                if 'file_path' in tool_output:
                    output += f"\n  {self._get_color('green')}📁 File:{self._get_color('end')} {tool_output['file_path']}"
                if 'url' in tool_output:
                    output += f"\n  {self._get_color('blue')}🌐 URL:{self._get_color('end')} {tool_output['url']}"
                return output
            else:
                return self._indent_text(json.dumps(tool_output, indent=2))
        elif isinstance(tool_output, str):
            return self._indent_text(tool_output)
        else:
            return self._indent_text(str(tool_output))

    def _format_text_content(self, text: str) -> str:
        """Format text content with proper indentation (backward compatibility)"""
        return self._format_enhanced_text_content(text)

    def _format_enhanced_text_content(self, text: str) -> str:
        """Enhanced text content formatting with better typography and code highlighting"""
        if not text:
            return f"  {self._get_color('grey')}(No content){self._get_color('end')}"
        
        # Check if text contains code blocks
        if '```' in text:
            return self._format_text_with_code_blocks(text)
        
        # Split into paragraphs and format
        paragraphs = text.split('\n\n')
        formatted_paragraphs = []
        
        for para in paragraphs:
            if para.strip():
                # Wrap long lines
                wrapped = self._wrap_text(para.strip(), 80)
                formatted_paragraphs.append(self._indent_text(wrapped))
        
        return "\n\n".join(formatted_paragraphs)

    def _format_text_with_code_blocks(self, text: str) -> str:
        """Format text with enhanced code block display"""
        lines = text.split('\n')
        formatted_lines = []
        in_code_block = False
        code_language = ""
        code_lines = []
        
        for line in lines:
            if line.strip().startswith('```'):
                if not in_code_block:
                    # Starting code block
                    in_code_block = True
                    code_language = line.strip()[3:].strip()  # Extract language
                    code_lines = []
                else:
                    # Ending code block
                    in_code_block = False
                    # Format the collected code block
                    formatted_code = self._format_code_block(code_lines, code_language)
                    formatted_lines.append(formatted_code)
                    code_language = ""
            elif in_code_block:
                code_lines.append(line)
            else:
                # Regular text line
                if line.strip():
                    wrapped = self._wrap_text(line.strip(), 80)
                    formatted_lines.append(self._indent_text(wrapped))
                else:
                    formatted_lines.append("")
        
        return "\n".join(formatted_lines)

    def _format_code_block(self, code_lines: List[str], language: str) -> str:
        """Format a code block with syntax highlighting and borders"""
        if not code_lines:
            return ""
        
        # Remove empty lines at start and end
        while code_lines and not code_lines[0].strip():
            code_lines.pop(0)
        while code_lines and not code_lines[-1].strip():
            code_lines.pop()
        
        if not code_lines:
            return ""
        
        # Calculate max width for the code block
        max_width = max(len(line) for line in code_lines)
        block_width = min(max_width + 4, 76)  # Leave some margin
        
        # Language indicator
        lang_indicator = f"  {self._get_color('yellow')}📝 {language.upper() if language else 'CODE'}{self._get_color('end')}"
        
        # Top border
        top_border = f"  {self._get_color('cyan')}{'┌' + '─' * (block_width - 2) + '┐'}{self._get_color('end')}"
        
        # Code lines with syntax highlighting
        formatted_code_lines = []
        for line in code_lines:
            # Basic syntax highlighting for common patterns
            highlighted_line = self._highlight_syntax(line, language)
            # Pad the line to match block width
            padded_line = highlighted_line.ljust(block_width - 2)
            formatted_code_lines.append(f"  {self._get_color('cyan')}│{self._get_color('end')} {padded_line} {self._get_color('cyan')}│{self._get_color('end')}")
        
        # Bottom border
        bottom_border = f"  {self._get_color('cyan')}{'└' + '─' * (block_width - 2) + '┘'}{self._get_color('end')}"
        
        return f"{lang_indicator}\n{top_border}\n" + "\n".join(formatted_code_lines) + f"\n{bottom_border}"

    def _highlight_syntax(self, line: str, language: str) -> str:
        """Basic syntax highlighting for code lines"""
        if not self.supports_colors:
            return line
        
        # JavaScript/JSX highlighting
        if language.lower() in ['js', 'jsx', 'javascript', 'react']:
            # Keywords
            keywords = ['const', 'let', 'var', 'function', 'return', 'if', 'else', 'for', 'while', 'switch', 'case', 'default', 'import', 'export', 'from', 'useState', 'useEffect']
            # Strings
            line = self._highlight_pattern(line, r'"[^"]*"', 'green')
            line = self._highlight_pattern(line, r"'[^']*'", 'green')
            # Keywords
            for keyword in keywords:
                line = line.replace(keyword, f"{self._get_color('purple')}{keyword}{self._get_color('end')}")
            # Numbers
            line = self._highlight_pattern(line, r'\b\d+\b', 'yellow')
            # Comments
            line = self._highlight_pattern(line, r'//.*$', 'grey')
        
        # Python highlighting
        elif language.lower() in ['py', 'python']:
            # Keywords
            keywords = ['def', 'class', 'import', 'from', 'return', 'if', 'else', 'elif', 'for', 'while', 'try', 'except', 'finally', 'with', 'as', 'True', 'False', 'None']
            # Strings
            line = self._highlight_pattern(line, r'"[^"]*"', 'green')
            line = self._highlight_pattern(line, r"'[^']*'", 'green')
            # Keywords
            for keyword in keywords:
                line = line.replace(keyword, f"{self._get_color('purple')}{keyword}{self._get_color('end')}")
            # Numbers
            line = self._highlight_pattern(line, r'\b\d+\b', 'yellow')
            # Comments
            line = self._highlight_pattern(line, r'#.*$', 'grey')
        
        # CSS highlighting
        elif language.lower() in ['css']:
            # Properties
            line = self._highlight_pattern(line, r'[a-zA-Z-]+(?=\s*:)', 'cyan')
            # Values
            line = self._highlight_pattern(line, r':\s*[^;]+', 'yellow')
            # Comments
            line = self._highlight_pattern(line, r'/\*.*?\*/', 'grey')
        
        return line

    def _highlight_pattern(self, text: str, pattern: str, color: str) -> str:
        """Highlight text matching a regex pattern with color"""
        import re
        matches = re.finditer(pattern, text)
        # Process matches in reverse order to avoid index shifting
        for match in reversed(list(matches)):
            start, end = match.span()
            matched_text = text[start:end]
            colored_text = f"{self._get_color(color)}{matched_text}{self._get_color('end')}"
            text = text[:start] + colored_text + text[end:]
        return text

    def _format_final_result(self, result: str) -> str:
        """Format final result with highlight box and code detection"""
        if not result:
            return f"  {self._get_color('grey')}(No result){self._get_color('end')}"
        
        result_str = str(result)
        
        # Check if result contains code-like content
        if self._looks_like_code(result_str):
            return self._format_code_result(result_str)
        
        # Regular text result
        result_lines = result_str.split('\n')
        max_width = max(len(line) for line in result_lines) if result_lines else 0
        box_width = min(max_width + 4, 76)  # Leave some margin
        
        output = []
        output.append(f"  {self._get_color('orange')}{'┌' + '─' * box_width + '┐'}{self._get_color('end')}")
        
        for line in result_lines:
            padded_line = line.ljust(box_width)
            output.append(f"  {self._get_color('orange')}│{self._get_color('end')} {padded_line} {self._get_color('orange')}│{self._get_color('end')}")
        
        output.append(f"  {self._get_color('orange')}{'└' + '─' * box_width + '┘'}{self._get_color('end')}")
        
        return "\n".join(output)

    def _looks_like_code(self, text: str) -> bool:
        """Detect if text looks like code"""
        code_indicators = [
            'function', 'const', 'let', 'var', 'def', 'class', 'import', 'export',
            'return', 'if', 'else', 'for', 'while', 'try', 'catch', '=>', '()', '{}',
            'console.log', 'print(', 'import ', 'from ', 'require('
        ]
        
        text_lower = text.lower()
        code_score = sum(1 for indicator in code_indicators if indicator in text_lower)
        
        # Also check for common code patterns
        has_brackets = text.count('{') > 0 or text.count('(') > 0
        has_semicolons = text.count(';') > 0
        has_equals = text.count('=') > 0
        
        return code_score >= 2 or (has_brackets and has_semicolons) or (has_brackets and has_equals)

    def _format_code_result(self, code: str) -> str:
        """Format code result with syntax highlighting"""
        lines = code.split('\n')
        
        # Detect language based on content
        language = self._detect_language(code)
        
        # Remove empty lines at start and end
        while lines and not lines[0].strip():
            lines.pop(0)
        while lines and not lines[-1].strip():
            lines.pop()
        
        if not lines:
            return f"  {self._get_color('grey')}(Empty code){self._get_color('end')}"
        
        # Calculate max width
        max_width = max(len(line) for line in lines)
        block_width = min(max_width + 4, 76)
        
        # Language indicator
        lang_indicator = f"  {self._get_color('yellow')}💻 {language.upper()} OUTPUT{self._get_color('end')}"
        
        # Top border
        top_border = f"  {self._get_color('green')}{'┌' + '─' * (block_width - 2) + '┐'}{self._get_color('end')}"
        
        # Code lines with syntax highlighting
        formatted_lines = []
        for line in lines:
            highlighted_line = self._highlight_syntax(line, language)
            padded_line = highlighted_line.ljust(block_width - 2)
            formatted_lines.append(f"  {self._get_color('green')}│{self._get_color('end')} {padded_line} {self._get_color('green')}│{self._get_color('end')}")
        
        # Bottom border
        bottom_border = f"  {self._get_color('green')}{'└' + '─' * (block_width - 2) + '┘'}{self._get_color('end')}"
        
        return f"{lang_indicator}\n{top_border}\n" + "\n".join(formatted_lines) + f"\n{bottom_border}"

    def _detect_language(self, code: str) -> str:
        """Detect programming language from code content"""
        code_lower = code.lower()
        
        if any(keyword in code_lower for keyword in ['function', 'const', 'let', 'var', 'console.log', '=>']):
            return 'javascript'
        elif any(keyword in code_lower for keyword in ['def', 'import ', 'print(', 'True', 'False']):
            return 'python'
        elif any(keyword in code_lower for keyword in ['class', 'public', 'private', 'void', 'System.out']):
            return 'java'
        elif any(keyword in code_lower for keyword in ['#include', 'int main', 'printf', 'scanf']):
            return 'c'
        elif any(keyword in code_lower for keyword in ['package', 'import ', 'func ', 'fmt.']):
            return 'go'
        elif any(keyword in code_lower for keyword in ['fn ', 'let ', 'println!', 'use ']):
            return 'rust'
        else:
            return 'code'

    def _format_footer(self) -> str:
        """Format footer (backward compatibility)"""
        return self._format_enhanced_footer()

    def _format_enhanced_footer(self) -> str:
        """Enhanced footer with better styling"""
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        border = f"{self._get_color('cyan')}{'═' * 60}{self._get_color('end')}"
        timestamp_line = f"{self._get_color('grey')}{'Generated at ' + timestamp:^60}{self._get_color('end')}"
        return f"{border}\n{timestamp_line}\n{border}"

    # Keep existing methods for compatibility
    def _colorize_diff(self, diff_text: str) -> str:
        """Colorize unified diff output for terminal display."""
        lines = diff_text.split('\n')
        colored_lines = []
        for line in lines:
            if line.startswith('+') and not line.startswith('+++'):
                colored_lines.append(f"{self._get_color('green')}{line}{self._get_color('end')}")
            elif line.startswith('-') and not line.startswith('---'):
                colored_lines.append(f"{self._get_color('red')}{line}{self._get_color('end')}")
            elif line.startswith('@@'):
                colored_lines.append(f"{self._get_color('yellow')}{line}{self._get_color('end')}")
            elif line.startswith('+++') or line.startswith('---'):
                colored_lines.append(f"{self._get_color('yellow')}{line}{self._get_color('end')}")
            elif line.strip() == '' or line.startswith(' '):
                colored_lines.append(f"{self._get_color('grey')}{line}{self._get_color('end')}")
            else:
                colored_lines.append(line)
        return '\n'.join(colored_lines)
    
    def _format_extracted_content(self, content: Dict[str, Any]) -> str:
        """Format extracted content from web scraping"""
        output = []
        
        if 'title' in content and content['title']:
            output.append(f"    {self._get_color('blue')}📰 Title:{self._get_color('end')} {content['title']}")
        
        if 'text' in content and content['text']:
            text_preview = content['text'][:300] + "..." if len(content['text']) > 300 else content['text']
            output.append(f"    {self._get_color('cyan')}📝 Text Preview:{self._get_color('end')}\n{self._indent_text(text_preview, 6)}")
        
        if 'links' in content and content['links']:
            output.append(f"    {self._get_color('green')}🔗 Links Found:{self._get_color('end')} {len(content['links'])}")
            for i, link in enumerate(content['links'][:3]):
                output.append(f"      {i+1}. {link.get('text', 'No text')[:50]}...")
        
        if 'images' in content and content['images']:
            output.append(f"    {self._get_color('purple')}🖼️  Images Found:{self._get_color('end')} {len(content['images'])}")
        
        return "\n".join(output)
    
    def _indent_text(self, text: str, indent: int = 2) -> str:
        """Indent text by specified number of spaces"""
        indent_str = " " * indent
        return "\n".join(indent_str + line for line in text.split('\n'))
    
    def _wrap_text(self, text: str, width: int = 80) -> str:
        """Wrap text to specified width"""
        import textwrap
        return textwrap.fill(text, width=width)
    
    def _extract_final_result(self, tool_output: Any) -> str:
        """Extract the most relevant raw result from the tool output for the Final Result section."""
        # Prefer executed code result if present
        if isinstance(tool_output, dict) and 'python_tool_execution' in tool_output:
            exec_result = tool_output['python_tool_execution']
            if isinstance(exec_result, dict):
                # Show stdout if successful, else stderr
                if exec_result.get('success'):
                    out = exec_result.get('stdout', '').strip()
                    if out:
                        return out
                err = exec_result.get('stderr', '').strip()
                if err:
                    return err
        # If tool_output is a dict with one key, return its value if it's a primitive or string
        if isinstance(tool_output, dict):
            for v in tool_output.values():
                if isinstance(v, (str, int, float, bool)):
                    return v
                if isinstance(v, dict):
                    if 'result' in v:
                        return v['result']
                    if 'output' in v:
                        return v['output']
        if isinstance(tool_output, (str, int, float, bool)):
            return tool_output
        return None

    # Enhanced versions of existing methods
    def format_error(self, error_message: str, context: Optional[Dict] = None) -> str:
        """Enhanced error message formatting"""
        output = []
        
        output.append(self._format_enhanced_header("❌ Error"))
        output.append("")
        output.append(f"{self._get_color('red')}🚨 Error Message:{self._get_color('end')}")
        output.append(self._indent_text(error_message))
        
        if context:
            output.append("")
            output.append(f"{self._get_color('yellow')}📋 Context:{self._get_color('end')}")
            output.append(self._indent_text(json.dumps(context, indent=2)))
        
        output.append("")
        output.append(self._format_enhanced_footer())
        
        return "\n".join(output)
    
    def format_memory_stats(self, stats: Dict[str, Any]) -> str:
        """Enhanced memory statistics formatting"""
        output = []
        
        output.append(self._format_enhanced_header("💾 Memory Statistics"))
        output.append("")
        
        for key, value in stats.items():
            formatted_key = key.replace('_', ' ').title()
            output.append(f"  {self._get_color('blue')}📊{self._get_color('end')} {formatted_key}: {self._get_color('green')}{value}{self._get_color('end')}")
        
        output.append("")
        output.append(self._format_enhanced_footer())
        
        return "\n".join(output)
    
    def format_tool_list(self, tools: Dict[str, Any]) -> str:
        """Enhanced tool list formatting"""
        output = []
        
        output.append(self._format_enhanced_header("🔧 Available Tools"))
        output.append("")
        
        for tool_name, tool_info in tools.items():
            output.append(f"{self._get_color('bold')}{self._get_color('green')}🛠️  {tool_name}{self._get_color('end')}")
            output.append(f"    {self._get_color('blue')}📝 Description:{self._get_color('end')} {tool_info.get('description', 'No description')}")
            
            capabilities = tool_info.get('capabilities', [])
            if capabilities:
                output.append(f"    {self._get_color('yellow')}⚡ Capabilities:{self._get_color('end')} {', '.join(capabilities)}")
            
            output.append("")
        
        output.append(self._format_enhanced_footer())
        
        return "\n".join(output)
    
    def format_simple_message(self, message: str, message_type: str = "info") -> str:
        """Enhanced simple message formatting"""
        icons = {
            "info": "ℹ️",
            "success": "✅",
            "warning": "⚠️",
            "error": "❌"
        }
        
        colors = {
            "info": self._get_color('blue'),
            "success": self._get_color('green'),
            "warning": self._get_color('yellow'),
            "error": self._get_color('red')
        }
        
        icon = icons.get(message_type, "ℹ️")
        color = colors.get(message_type, self._get_color('blue'))
        
        return f"{color}{icon} {message}{self._get_color('end')}"


# Keep the original OutputFormatter for backward compatibility
class OutputFormatter(EnhancedOutputFormatter):
    """Backward compatibility wrapper for the original OutputFormatter"""
    pass

