tools/formatting_tools.py

"""
Formatting tools for displaying analysis results with rich formatting.
"""

import json
from loguru import logger
from smolagents import Tool
from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from rich.table import Table
from rich.markdown import Markdown
from rich.theme import Theme
from rich.box import ROUNDED
from rich.console import Group


class FormatAnalysisResultsTool(Tool):
    """Tool for formatting analysis results"""
    
    name = "format_analysis_results"
    description = "Formats the analysis results for better readability"
    inputs = {
        "analysis_json": {"type": "any", "description": "JSON string or dictionary containing the analysis results"},
        "pretty": {"type": "boolean", "description": "Whether to use rich formatting (True) or simple text formatting (False)", "nullable": True}
    }
    output_type = "string"
    
    def forward(self, analysis_json, pretty: bool = True) -> str:
        """
        Formats the analysis results for better readability.
        
        Args:
            analysis_json: JSON string or dictionary containing the analysis results
            pretty: Whether to use rich formatting (True) or simple text formatting (False)
            
        Returns:
            A formatted string representation of the analysis
        """
        try:
            # Parse the JSON string into a Python dictionary if it's a string
            if isinstance(analysis_json, str):
                analysis = json.loads(analysis_json)
            else:
                analysis = analysis_json
                
            if pretty:
                # Rich formatting with the rich library
                try:
                    # Create a string buffer to capture the rich output
                    console = Console(record=True, width=100)
                    
                    # Create a custom theme for consistent styling
                    # Using direct style definitions instead of named styles
                    custom_theme = Theme({
                        "heading": "bold cyan underline",
                        "highlight": "bold yellow",
                        "positive": "green",
                        "negative": "red",
                        "neutral": "magenta",
                        "quote": "italic yellow",
                        "metadata": "dim white"
                    })
                    
                    # Apply the theme to our console
                    console.theme = custom_theme
                    
                    # Format the summary section
                    summary = analysis.get("summary", "No summary available")
                    console.print(Panel(Text(summary, justify="center"), title="[heading]Song Analysis[/]", subtitle="[metadata]Summary[/]"))
                    
                    # Create a table for the main themes and mood
                    info_table = Table(show_header=False, box=ROUNDED, expand=True)
                    info_table.add_column("Key", style="bold blue")
                    info_table.add_column("Value")
                    
                    # Add the mood
                    mood = analysis.get("mood", "Not specified")
                    info_table.add_row("Mood", mood)
                    
                    # Add the themes as a comma-separated list
                    themes = analysis.get("main_themes", [])
                    if themes:
                        themes_text = ", ".join([f"[highlight]{theme}[/]" for theme in themes])
                        info_table.add_row("Main Themes", themes_text)
                    
                    console.print(info_table)
                    
                    # Section-by-section analysis
                    sections = analysis.get("sections_analysis", [])
                    if sections:
                        console.print("\n[heading]Section-by-Section Analysis[/]")
                        
                        for section in sections:
                            section_type = section.get("section_type", "Unknown")
                            section_number = section.get("section_number", "")
                            section_title = f"{section_type.title()} {section_number}" if section_number else section_type.title()
                            
                            section_analysis = section.get("analysis", "No analysis available")
                            lines = section.get("lines", [])
                            
                            # Create a group with the lyrics and analysis
                            section_content = []
                            
                            if lines:
                                lyrics_text = "\n".join([f"> [quote]{line}[/]" for line in lines])
                                section_content.append(Text("Lyrics:", style="bold blue"))
                                section_content.append(Markdown(lyrics_text))
                            
                            section_content.append(Text("Analysis:", style="bold blue"))
                            section_content.append(Text(section_analysis))
                            
                            # Add the section panel
                            console.print(Panel(
                                Group(*section_content),
                                title=f"[bold cyan]{section_title}[/]",
                                border_style="cyan"
                            ))
                    
                    # Significant lines
                    sig_lines = analysis.get("significant_lines", [])
                    if sig_lines:
                        console.print("\n[heading]Significant Lines[/]")
                        
                        for i, line_data in enumerate(sig_lines):
                            line = line_data.get("line", "")
                            significance = line_data.get("significance", "")
                            
                            console.print(Panel(
                                f"[quote]\"{line}\"[/]\n\n[bold blue]Significance:[/] {significance}",
                                title=f"[highlight]Key Line #{i+1}[/]",
                                border_style="highlight"
                            ))
                    
                    # Conclusion
                    conclusion = analysis.get("conclusion", "No conclusion available")
                    console.print("\n[heading]Conclusion[/]")
                    console.print(Panel(conclusion, border_style="neutral"))
                    
                    # Export the rich text as a string
                    return console.export_text()
                    
                except Exception as e:
                    logger.error("Error in rich formatting: {}", str(e))
                    # Fall back to simple formatting if rich formatting fails
                    logger.info("Falling back to simple text formatting")
                    pretty = False
            
            if not pretty:
                # Simple text formatting
                formatted_text = []
                
                # Summary
                formatted_text.append("SONG ANALYSIS SUMMARY")
                formatted_text.append("====================")
                formatted_text.append(analysis.get("summary", "No summary available"))
                formatted_text.append("")
                
                # Main themes
                themes = analysis.get("main_themes", [])
                if themes:
                    formatted_text.append("MAIN THEMES")
                    formatted_text.append("===========")
                    for theme in themes:
                        formatted_text.append(f"• {theme}")
                    formatted_text.append("")
                
                # Mood
                mood = analysis.get("mood", "Not specified")
                formatted_text.append("MOOD")
                formatted_text.append("====")
                formatted_text.append(mood)
                formatted_text.append("")
                
                # Sections analysis
                sections = analysis.get("sections_analysis", [])
                if sections:
                    formatted_text.append("SECTION-BY-SECTION ANALYSIS")
                    formatted_text.append("==========================")
                    
                    for i, section in enumerate(sections):
                        section_type = section.get("section_type", "Unknown")
                        section_number = section.get("section_number", i+1)
                        section_analysis = section.get("analysis", "No analysis available")
                        
                        formatted_text.append(f"{section_type.upper()} {section_number}")
                        formatted_text.append("-" * (len(section_type) + len(str(section_number)) + 1))
                        
                        # Format the section lines
                        lines = section.get("lines", [])
                        if lines:
                            formatted_text.append("Lyrics:")
                            for line in lines:
                                formatted_text.append(f"> {line}")
                            formatted_text.append("")
                        
                        formatted_text.append("Analysis:")
                        formatted_text.append(section_analysis)
                        formatted_text.append("")
                
                # Significant lines
                sig_lines = analysis.get("significant_lines", [])
                if sig_lines:
                    formatted_text.append("SIGNIFICANT LINES")
                    formatted_text.append("=================")
                    
                    for i, line_data in enumerate(sig_lines):
                        line = line_data.get("line", "")
                        significance = line_data.get("significance", "")
                        
                        formatted_text.append(f"Key Line #{i+1}:")
                        formatted_text.append(f'"{line}"')
                        formatted_text.append(f"Significance: {significance}")
                        formatted_text.append("")
                
                # Conclusion
                conclusion = analysis.get("conclusion", "No conclusion available")
                formatted_text.append("CONCLUSION")
                formatted_text.append("==========")
                formatted_text.append(conclusion)
                
                return "\n".join(formatted_text)
                
        except json.JSONDecodeError:
            logger.error("Failed to parse analysis JSON: {}", analysis_json[:100] + "..." if len(analysis_json) > 100 else analysis_json)
            return f"Error: Could not parse the analysis result as JSON. Raw output:\n{analysis_json}"
        except Exception as e:
            logger.error("Error formatting analysis: {}", str(e))
            return f"Error formatting analysis: {str(e)}\nRaw output:\n{analysis_json}"