from typing import Any, Optional, Dict, Union
from smolagents.tools import Tool

class FinalAnswerTool(Tool):
    name = "final_answer"
    description = "Provides a final answer to the given problem with proper formatting and sources."
    inputs = {
        'content': {'type': 'string', 'description': 'The main content of the answer'},
        'sources': {
            'type': 'array', 
            'description': 'List of sources used in the research',
            'nullable': True
        },
        'metadata': {
            'type': 'object', 
            'description': 'Additional metadata like readability score, SEO info, etc.',
            'nullable': True
        }
    }
    output_type = "string"

    def format_output(self, content: str, sources: list = None, metadata: Dict = None) -> str:
        """Format the output in a clean, readable way with proper sections."""
        output_parts = []
        
        # Add main content
        output_parts.append("# Content\n")
        output_parts.append(content.strip())
        output_parts.append("\n\n")
        
        # Add sources if available
        if sources and len(sources) > 0:
            output_parts.append("# Sources\n")
            for idx, source in enumerate(sources, 1):
                output_parts.append(f"{idx}. {source}\n")
            output_parts.append("\n")
        
        # Add metadata if available
        if metadata:
            output_parts.append("# Additional Information\n")
            for key, value in metadata.items():
                output_parts.append(f"- {key}: {value}\n")
        
        return "".join(output_parts)

    def forward(self, content: str, sources: list = None, metadata: Dict = None) -> str:
        """Process the final answer with proper formatting."""
        return self.format_output(content, sources, metadata)

    def __init__(self, *args, **kwargs):
        self.is_initialized = False