import os
import datetime
import requests
import pytz
import yaml

# Try different import patterns for smolagents
try:
    from smolagents import CodeAgent, DuckDuckGoSearchTool, load_tool, tool
    print("✓ Core smolagents imports successful")
except ImportError as e:
    print(f"✗ Error importing from smolagents: {e}")
    raise

# Try to import the model class with different names
model = None
try:
    from smolagents import HfApiModel
    model_class = HfApiModel
    print("✓ HfApiModel imported")
except ImportError:
    try:
        from smolagents.models import HfApiModel
        model_class = HfApiModel
        print("✓ HfApiModel imported from models")
    except ImportError:
        try:
            from smolagents import HuggingFaceModel
            model_class = HuggingFaceModel
            print("✓ HuggingFaceModel imported")
        except ImportError:
            try:
                from smolagents.models import HuggingFaceModel
                model_class = HuggingFaceModel
                print("✓ HuggingFaceModel imported from models")
            except ImportError:
                print("✗ No suitable model class found, will use default")
                model_class = None

# Try importing final answer tool
try:
    from tools.final_answer import FinalAnswerTool
    print("✓ FinalAnswerTool imported")
except ImportError as e:
    print(f"! FinalAnswerTool not found: {e}")
    # Create a simple fallback
    @tool
    def final_answer(answer: str) -> str:
        """Provide the final answer to the user."""
        return answer
    FinalAnswerTool = None

# Try importing Gradio UI
try:
    from Gradio_UI import GradioUI
    print("✓ GradioUI imported")
except ImportError as e:
    print(f"! GradioUI not found: {e}")
    GradioUI = None

# -------------------------
# Local custom tools
# -------------------------
@tool
def my_custom_tool(arg1: str, arg2: int) -> str:
    """A tool that does nothing yet 
    Args:
        arg1: the first argument
        arg2: the second argument
    """
    return "What magic will you build ?"

@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:
        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)}"

@tool
def calculator(a: int, b: int) -> int:
    """Multiply two integers.
    Args:
        a: first integer
        b: second integer
    """
    return int(a) * int(b)

# Simple weather tool using Open-Meteo (no key needed)
CITY_COORDS = {
    "new york": (40.7128, -74.0060),
    "paris": (48.8566, 2.3522),
    "london": (51.5074, -0.1278),
    "tokyo": (35.6762, 139.6503),
}

@tool
def get_weather(city: str) -> str:
    """Get current temperature and humidity for a known city (New York, Paris, London, Tokyo).
    Args:
        city: City name (case-insensitive)
    """
    coords = CITY_COORDS.get(city.lower())
    if not coords:
        return f"Unknown city '{city}'. Try one of: {', '.join(CITY_COORDS.keys())}."
    lat, lon = coords
    url = (
        "https://api.open-meteo.com/v1/forecast"
        f"?latitude={lat}&longitude={lon}"
        "&current=temperature_2m,relative_humidity_2m,weather_code"
    )
    try:
        r = requests.get(url, timeout=10)
        r.raise_for_status()
        cur = r.json().get("current", {})
        t = cur.get("temperature_2m")
        h = cur.get("relative_humidity_2m")
        if t is None or h is None:
            return "Weather service did not return expected fields."
        return f"Current weather in {city.title()}: {t}°C, {h}% humidity."
    except Exception as e:
        return f"Weather fetch failed: {e}"

# -------------------------
# Simple image generation tool (fallback)
# -------------------------
@tool
def simple_image_description(prompt: str) -> str:
    """Generate a detailed description of an image based on the prompt.
    This is a fallback when actual image generation is not available.
    Args:
        prompt: Description of the image to generate
    """
    return f"Image description: {prompt}. Note: This is a text description as image generation is currently unavailable. To enable actual image generation, please add your Hugging Face token to the space settings."

# -------------------------
# Required final answer tool
# -------------------------
if FinalAnswerTool:
    try:
        final_answer_tool = FinalAnswerTool()
        print("✓ Final answer tool loaded")
    except Exception as e:
        print(f"✗ Error loading final answer tool: {e}")
        final_answer_tool = final_answer
else:
    final_answer_tool = final_answer

# -------------------------
# Model setup with error handling
# -------------------------
if model_class:
    try:
        model = model_class(
            max_tokens=2096,
            temperature=0.5,
            model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
        )
        print("✓ Model loaded successfully")
    except Exception as e:
        print(f"✗ Error loading model: {e}")
        # Try a smaller/simpler model
        try:
            model = model_class(
                max_tokens=1024,
                temperature=0.5,
                model_id='microsoft/DialoGPT-medium',
            )
            print("✓ Fallback model loaded")
        except Exception as e2:
            print(f"✗ Fallback model also failed: {e2}")
            # Try without specific model
            try:
                model = model_class()
                print("✓ Default model loaded")
            except Exception as e3:
                print(f"✗ Default model failed: {e3}")
                model = None
else:
    print("! No model class available - will try to create agent without explicit model")
    model = None

# -------------------------
# Pre-made tools from Hub / smolagents
# -------------------------
# DuckDuckGo web search (no key) - with error handling
try:
    duck_search = DuckDuckGoSearchTool()
    print("✓ DuckDuckGo search tool loaded")
except Exception as e:
    print(f"✗ Error loading DuckDuckGo tool: {e}")
    duck_search = None

# Text-to-image tool - with proper authentication handling
image_generation_tool = None
hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")

if hf_token:
    try:
        # Set the token for huggingface_hub
        from huggingface_hub import login
        login(token=hf_token)
        
        image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
        print("✓ Image generation tool loaded successfully")
    except Exception as e:
        print(f"✗ Failed to load image generation tool: {e}")
        print("  Using fallback image description tool instead")
        image_generation_tool = simple_image_description
else:
    print("! No HF_TOKEN found - using fallback image description tool")
    print("  To enable real image generation, add your Hugging Face token in Space settings")
    image_generation_tool = simple_image_description

# -------------------------
# Load prompts with error handling
# -------------------------
prompt_templates = {}
try:
    if os.path.exists("prompts.yaml"):
        with open("prompts.yaml", 'r') as stream:
            prompt_templates = yaml.safe_load(stream)
        print("✓ Prompts loaded from prompts.yaml")
    else:
        print("! prompts.yaml not found, using default prompts")
        prompt_templates = {
            "system_prompt": """You are a helpful AI assistant with access to various tools.
            
Available tools:
- calculator: For mathematical operations
- get_weather: For weather information in major cities
- get_current_time_in_timezone: For time information in different timezones
- duck_search: For web searches (if available)
- simple_image_description or image_generation_tool: For image-related requests
- final_answer: To provide your final response

When users request image generation:
1. If image_generation_tool is available, use it to create actual images
2. If only simple_image_description is available, provide detailed descriptions

Always try to use the most appropriate tool for the task and provide helpful responses."""
        }
except Exception as e:
    print(f"✗ Error loading prompts: {e}")
    prompt_templates = {}

# -------------------------
# Build tools list
# -------------------------
tools_list = [
    final_answer_tool,
    calculator,
    get_current_time_in_timezone,
    get_weather,
    my_custom_tool,
    image_generation_tool,  # This will be either the real tool or the fallback
]

if duck_search:
    tools_list.append(duck_search)

print(f"✓ Loaded {len(tools_list)} tools")

# -------------------------
# Agent with tools registered
# -------------------------
try:
    if model:
        agent = CodeAgent(
            model=model,
            tools=tools_list,
            max_steps=10,
            verbosity_level=2,
            prompt_templates=prompt_templates
        )
    else:
        # Try creating agent without explicit model
        agent = CodeAgent(
            tools=tools_list,
            max_steps=10,
            verbosity_level=2,
            prompt_templates=prompt_templates
        )
    print("✓ Agent created successfully")
except Exception as e:
    print(f"✗ Error creating agent: {e}")
    raise e

# -------------------------
# Launch UI
# -------------------------
if GradioUI:
    try:
        ui = GradioUI(agent)
        print("✓ UI created successfully")
        ui.launch()
    except Exception as e:
        print(f"✗ Error launching custom UI: {e}")
        GradioUI = None

if not GradioUI:
    # Fallback: create a simple Gradio interface
    print("Using fallback Gradio interface")
    import gradio as gr
    
    def chat_function(message):
        try:
            result = agent.run(message)
            return str(result)
        except Exception as e:
            return f"Error: {str(e)}"
    
    demo = gr.Interface(
        fn=chat_function,
        inputs=gr.Textbox(label="Your message", placeholder="Ask me anything..."),
        outputs=gr.Textbox(label="Agent response"),
        title="AI Agent",
        description="Chat with an AI agent that has access to various tools including weather, time, calculator, and web search."
    )
    demo.launch(server_name="0.0.0.0", server_port=7860)