from typing import Any, Dict, Union, List
from smolagents.tools import Tool
import logging
from datetime import datetime

logger = logging.getLogger(__name__)

class FinalAnswerTool(Tool):
    name = "final_answer"
    description = "Formats and returns the final research report"
    inputs = {
        'answer': {
            'type': 'any',
            'description': 'The final research report content'
        }
    }
    output_type = "string"  # Changed from 'str' to 'string'

    def __init__(self):
        super().__init__()
        self.sections = [
            "Executive Summary",
            "Introduction",
            "Methodology",
            "Findings",
            "Discussion",
            "Conclusion",
            "References"
        ]

    def _format_list_content(self, content: List) -> str:
        """Format list content into readable text"""
        try:
            return "\n".join([f"- {item}" for item in content])
        except Exception as e:
            logger.error(f"Error formatting list content: {e}")
            return str(content)

    def _format_dict_content(self, content: Dict) -> str:
        """Format dictionary content into sections"""
        formatted_sections = []
        
        for section in self.sections:
            section_content = content.get(section, f"{section} section was not provided")
            if isinstance(section_content, list):
                section_content = self._format_list_content(section_content)
            formatted_sections.extend([
                f"## {section}",
                str(section_content),
                ""  # Empty line for better readability
            ])
        
        return "\n".join(formatted_sections)

    def _format_report(self, content: Any) -> str:
        """Format content as a proper research report"""
        try:
            # Create report header
            report_parts = [
                "# Науковий звіт",
                f"*Згенеровано: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*\n"
            ]

            # Handle different content types
            if isinstance(content, dict):
                report_parts.append(self._format_dict_content(content))
            elif isinstance(content, list):
                report_parts.append(self._format_list_content(content))
            else:
                report_parts.append(str(content))

            return "\n".join(report_parts)

        except Exception as e:
            error_msg = f"Error formatting report: {str(e)}"
            logger.error(error_msg)
            return f"# Error in Report Generation\n\n{error_msg}\n\nRaw content:\n{str(content)}"

    def forward(self, answer: Any) -> str:
        """Process and return the final answer"""
        logger.info("Formatting final research report")
        try:
            formatted_report = self._format_report(answer)
            logger.info("Research report formatted successfully")
            return formatted_report
        except Exception as e:
            error_msg = f"Error formatting research report: {str(e)}"
            logger.error(error_msg)
            return error_msg