from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """Get current time in specified timezone.
    Args:
        timezone: Valid timezone (e.g., 'America/New_York')
    """
    try:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"Current time in {timezone}: {local_time}"
    except Exception as e:
        return f"Error: {str(e)}"

# Performance optimizations
final_answer = FinalAnswerTool()

# Try alternative model endpoint if main is overloaded
model = HfApiModel(
    max_tokens=1024,  # Reduced from 2096 to limit output length
    temperature=0.3,  # Reduced from 0.5 for more focused responses
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
)

# Fallback model configuration
fallback_model = HfApiModel(
    max_tokens=1024,
    temperature=0.3,
    model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',
    custom_role_conversions=None,
)

# Load image generation tool (only if needed)
# image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

# Optimized prompt templates
optimized_prompts = {
    "system": """You are a helpful coding assistant. Be concise and direct.
    Focus on the specific task. Avoid unnecessary explanations unless requested.""",
    "user": """Task: {task}
    
    Available tools: {tools}
    
    Provide a clear, step-by-step solution."""
}

# Try to load custom prompts, fallback to optimized ones
try:
    with open("prompts.yaml", 'r') as stream:
        prompt_templates = yaml.safe_load(stream)
        # Add length limits to existing prompts
        if 'system' in prompt_templates:
            prompt_templates['system'] = prompt_templates['system'][:500]  # Limit system prompt
except FileNotFoundError:
    prompt_templates = optimized_prompts

# Create agent with performance optimizations
agent = CodeAgent(
    model=model,
    tools=[final_answer, get_current_time_in_timezone],
    max_steps=3,  # Reduced from 6 to limit processing time
    verbosity_level=0,  # Reduced from 1 to minimize output
    grammar=None,
    planning_interval=2,  # Add planning interval for better step management
    name="OptimizedAgent",
    description="Fast, efficient coding assistant",
    prompt_templates=prompt_templates
)

# Add error handling and timeout
def launch_with_optimizations():
    try:
        # Try main model first
        GradioUI(agent).launch(
            share=False,  # Disable sharing for better performance
            server_name="127.0.0.1",
            server_port=7860,
            max_threads=4  # Limit concurrent threads
        )
    except Exception as e:
        print(f"Main model failed: {e}")
        print("Switching to fallback model...")
        agent.model = fallback_model
        GradioUI(agent).launch()

if __name__ == "__main__":
    launch_with_optimizations()