from smolagents import CodeAgent, Tool, DuckDuckGoSearchTool, load_tool, tool
from models import OpenAIModel  # Import our local model
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
import re
import os
import logging
from jinja2 import Template, StrictUndefined

from Gradio_UI import GradioUI

# Configure logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

# Below is an example of a tool that does nothing. Amaze us with your creativity !
@tool
def calculate_min_price(prices: list[float])-> str: #it's import to specify the return type
    """A tool that calculates the min price from list of product prices
    Args:
        prices: list of product prices of
    """
    min_price =min(prices)
    return f"The minimum price is {min_price}"

@tool
def extract_price_from_snippet(snippet: str) -> list[str]:
    """
    A simple function to extract prices from a text snippet using regex.
    You can enhance this function for more complex price extraction.
    Args:
        snippet: text of all prices
    """
    # A basic regular expression to detect common price formats like $29.99, 29.99 USD, etc.
    price_pattern = r'\$\d+(?:,\d{3})*(?:\.\d{2})?|\d+(?:,\d{3})*(?:\.\d{2})?\s*(USD|EUR|GBP|INR|AUD|CAD)?'
    matches = re.findall(price_pattern, snippet)
    matches = [str(x) for x in matches]
    return matches


@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """A tool that fetches the current local time in a specified timezone.
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    """
    try:
        # Create timezone object
        tz = pytz.timezone(timezone)
        # Get current time in that timezone
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"


final_answer = FinalAnswerTool()

# Load configuration from agent.json
with open("agent.json", 'r') as f:
    config = yaml.safe_load(f)

# Initialize OpenAI model using the configuration
model = OpenAIModel(
    model=config["model"]["data"]["model_id"],
    max_tokens=config["model"]["data"]["max_tokens"],
    temperature=config["model"]["data"]["temperature"],
    api_key=os.getenv("OPENAI_API_KEY")
)

# Import tool from Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

# Load and validate templates from prompts.yaml
with open("prompts.yaml", 'r') as stream:
    yaml_content = yaml.safe_load(stream)
    if not isinstance(yaml_content, dict):
        raise ValueError("YAML content must be a dictionary")
    if 'prompt_templates' not in yaml_content:
        raise ValueError("YAML must contain 'prompt_templates' key")
    prompt_templates = yaml_content['prompt_templates']
    if not isinstance(prompt_templates, dict):
        raise ValueError("prompt_templates must be a dictionary")
    
    logger.debug("Starting template validation...")
    try:
        # Create test tools for validation
        class TestTool(Tool):
            name = "test_tool"
            description = "A test tool"
            inputs = {'input': {'type': 'string', 'description': 'Test input'}}
            output_type = "string"

            def forward(self, input: str) -> str:
                return "test output"

        test_tools = [TestTool()]  # Create a list of tools
        
        # Validate templates
        for key, value in prompt_templates.items():
            if isinstance(value, dict):
                for subkey, subvalue in value.items():
                    if isinstance(subvalue, str):
                        logger.debug(f"Validating template: {key}.{subkey}")
                        logger.debug(f"Template content: {subvalue[:200]}...")
                        
                        # Validate template
                        try:
                            # Create a template with at least one template node
                            template_str = subvalue
                            if '{{' not in template_str and '{%' not in template_str:
                                template_str = "{{ task }}\n" + template_str
                            
                            # Create a template with the task variable
                            template = Template(template_str, undefined=StrictUndefined)
                            rendered = template.render(
                                tools=test_tools,  # Pass list of tools
                                task="test",
                                name="test",
                                final_answer="test",
                                remaining_steps=1,
                                answer_facts="test"
                            )
                            logger.debug(f"Template render test successful for {key}.{subkey}")
                        except Exception as e:
                            logger.error(f"Template render test failed for {key}.{subkey}: {str(e)}")
                            raise
            elif isinstance(value, str):
                logger.debug(f"Validating template: {key}")
                logger.debug(f"Template content: {value[:200]}...")
                
                # Validate template
                try:
                    # Create a template with at least one template node
                    template_str = value
                    if '{{' not in template_str and '{%' not in template_str:
                        template_str = "{{ task }}\n" + template_str
                    
                    # Create a template with the task variable
                    template = Template(template_str, undefined=StrictUndefined)
                    rendered = template.render(
                        tools=test_tools,  # Pass list of tools
                        task="test",
                        name="test",
                        final_answer="test",
                        remaining_steps=1,
                        answer_facts="test"
                    )
                    logger.debug(f"Template render test successful for {key}")
                except Exception as e:
                    logger.error(f"Template render test failed for {key}: {str(e)}")
                    raise
        
        logger.debug("Template validation completed successfully")
        
    except Exception as e:
        logger.error(f"Error during template validation: {str(e)}")
        raise

# Create the agent with the templates
class CalculateMinPriceTool(Tool):
    name = "calculate_min_price"
    description = "Calculate the minimum price from a list of prices"
    inputs = {'prices': {'type': 'array', 'description': 'List of product prices (numbers)'}}
    output_type = "string"

    def forward(self, prices: list[float]) -> str:
        min_price = min(prices)
        return f"The minimum price is {min_price}"

class ExtractPriceFromSnippetTool(Tool):
    name = "extract_price_from_snippet"
    description = "Extract prices from a text snippet"
    inputs = {'snippet': {'type': 'string', 'description': 'Text containing prices'}}
    output_type = "array"

    def forward(self, snippet: str) -> list[str]:
        price_pattern = r'\$\d+(?:,\d{3})*(?:\.\d{2})?|\d+(?:,\d{3})*(?:\.\d{2})?\s*(USD|EUR|GBP|INR|AUD|CAD)?'
        matches = re.findall(price_pattern, snippet)
        matches = [str(x) for x in matches]
        return matches

class GetCurrentTimeInTimezoneTool(Tool):
    name = "get_current_time_in_timezone"
    description = "Get the current time in a specified timezone"
    inputs = {'timezone': {'type': 'string', 'description': 'A valid timezone (e.g., America/New_York)'}}
    output_type = "string"

    def forward(self, timezone: str) -> str:
        try:
            tz = pytz.timezone(timezone)
            local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
            return f"The current local time in {timezone} is: {local_time}"
        except Exception as e:
            return f"Error fetching time for timezone '{timezone}': {str(e)}"

# Create the agent with the templates
tools = [
    final_answer,  # This is already a Tool instance
    DuckDuckGoSearchTool(),  # This is already a Tool instance
    CalculateMinPriceTool(),  # This is a Tool subclass
    ExtractPriceFromSnippetTool(),  # This is a Tool subclass
    GetCurrentTimeInTimezoneTool()  # This is a Tool subclass
]

# Debug prints to inspect tools
logger.debug("Inspecting tools list:")
for i, tool in enumerate(tools):
    logger.debug(f"Tool {i}: {tool}")
    logger.debug(f"  Type: {type(tool)}")
    logger.debug(f"  Is Tool instance: {isinstance(tool, Tool)}")
    logger.debug(f"  Has name: {hasattr(tool, 'name')}")
    logger.debug(f"  Has description: {hasattr(tool, 'description')}")
    if hasattr(tool, 'name'):
        logger.debug(f"  Name: {tool.name}")
    if hasattr(tool, 'description'):
        logger.debug(f"  Description: {tool.description}")

# Verify all tools are proper Tool instances and have required attributes
for tool in tools:
    if not isinstance(tool, Tool):
        raise TypeError(f"Tool {tool} is not an instance of Tool or its subclass")
    if not hasattr(tool, 'name') or not hasattr(tool, 'description'):
        raise AttributeError(f"Tool {tool} is missing required attributes (name or description)")

# Create a dictionary of tools for template rendering
tools_dict = {tool.name: tool for tool in tools}
logger.debug("Created tools dictionary:")
for name, tool in tools_dict.items():
    logger.debug(f"  {name}: {tool}")

print("[DEBUG] Starting app.py initialization...")

# Create a custom CodeAgent class that handles both list and dictionary requirements
class CustomCodeAgent(CodeAgent):
    def __init__(self, *args, **kwargs):
        print("[DEBUG] Initializing CustomCodeAgent...")
        # Store the tools dictionary for template rendering
        self.tools_dict = kwargs.pop('tools_dict', {})
        # Initialize parent class first
        super().__init__(*args, **kwargs)

    def initialize_system_prompt(self):
        print("[DEBUG] Initializing system prompt...")
        # Override to use tools_dict for template rendering
        template = self.prompt_templates["system_prompt"]
        # Convert tools_dict to list for template rendering
        tools_list = list(self.tools_dict.values())
        # Use Jinja2 template rendering
        return Template(template, undefined=StrictUndefined).render(
            tools=tools_list,  # Pass tools as a list
            task=getattr(self, 'task', ''),
            managed_agents=getattr(self, 'managed_agents', []),
            authorized_imports=getattr(self, 'authorized_imports', [])
        )

print("[DEBUG] Creating agent instance...")
agent = CustomCodeAgent(
    model=model,
    tools=tools,  # Pass tools as a list for _setup_tools
    tools_dict=tools_dict,  # Pass tools as a dictionary for template rendering
    max_steps=15,
    verbosity_level=2,
    grammar=None,
    planning_interval=1,
    name="question_answering_agent",
    description="An agent specialized in answering various types of questions using available tools. The agent must use the final_answer tool to submit its answer.",
    prompt_templates=prompt_templates
)

print("[DEBUG] Building Gradio demo...")
demo = GradioUI(agent).build_blocks()
print("[DEBUG] Gradio demo built and exposed as 'demo'.")
print("[DEBUG] About to assign app = demo")
app = demo
print("[DEBUG] app assigned successfully")