from smolagents import CodeAgent, DuckDuckGoSearchTool, InferenceClientModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
import json
import os
import math
import re
from typing import Dict, List, Any, Optional
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# =============================================================================
# CUSTOM TOOLS COLLECTION
# =============================================================================

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """Get the current local time in a specified timezone.
    
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York', 'Europe/London', 'Asia/Tokyo')
    """
    try:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S %Z")
        return f"Current time in {timezone}: {local_time}"
    except Exception as e:
        return f"Error: Invalid timezone '{timezone}'. Please use format like 'America/New_York'"

@tool
def weather_forecast(city: str, country_code: str = "") -> str:
    """Get current weather information for a city.
    
    Args:
        city: Name of the city
        country_code: Optional 2-letter country code (e.g., 'US', 'UK', 'CA')
    """
    try:
        # Using OpenWeatherMap API (you'll need to get a free API key)
        api_key = os.getenv('OPENWEATHER_API_KEY', 'demo_key')
        
        if country_code:
            query = f"{city},{country_code}"
        else:
            query = city
            
        url = f"http://api.openweathermap.org/data/2.5/weather?q={query}&appid={api_key}&units=metric"
        
        if api_key == 'demo_key':
            return f"Weather service unavailable. Please set OPENWEATHER_API_KEY environment variable."
        
        response = requests.get(url, timeout=10)
        data = response.json()
        
        if response.status_code == 200:
            temp = data['main']['temp']
            feels_like = data['main']['feels_like']
            humidity = data['main']['humidity']
            description = data['weather'][0]['description']
            wind_speed = data['wind']['speed']
            
            return f"Weather in {city}: {description.title()}, {temp}°C (feels like {feels_like}°C), Humidity: {humidity}%, Wind: {wind_speed} m/s"
        else:
            return f"Could not fetch weather for {city}. Error: {data.get('message', 'Unknown error')}"
            
    except Exception as e:
        return f"Weather service error: {str(e)}"

@tool
def calculate_advanced_math(expression: str) -> str:
    """Safely evaluate mathematical expressions including advanced functions.
    
    Args:
        expression: Mathematical expression (e.g., 'sqrt(16)', 'sin(pi/2)', '2**3 + log(10)')
    """
    try:
        # Safe math evaluation with common functions
        safe_dict = {
            '__builtins__': {},
            'abs': abs, 'round': round, 'min': min, 'max': max,
            'sum': sum, 'pow': pow,
            'sqrt': math.sqrt, 'sin': math.sin, 'cos': math.cos, 'tan': math.tan,
            'log': math.log, 'log10': math.log10, 'exp': math.exp,
            'pi': math.pi, 'e': math.e,
            'ceil': math.ceil, 'floor': math.floor,
            'factorial': math.factorial,
            'degrees': math.degrees, 'radians': math.radians
        }
        
        result = eval(expression, safe_dict)
        return f"Result: {result}"
        
    except Exception as e:
        return f"Math error: {str(e)}. Please check your expression syntax."

@tool
def text_analyzer(text: str) -> str:
    """Analyze text for various metrics and insights.
    
    Args:
        text: The text to analyze
    """
    try:
        # Basic text statistics
        word_count = len(text.split())
        char_count = len(text)
        char_count_no_spaces = len(text.replace(' ', ''))
        sentence_count = len([s for s in re.split(r'[.!?]+', text) if s.strip()])
        paragraph_count = len([p for p in text.split('\n\n') if p.strip()])
        
        # Average metrics
        avg_words_per_sentence = word_count / max(sentence_count, 1)
        avg_chars_per_word = char_count_no_spaces / max(word_count, 1)
        
        # Reading time estimate (average 200 words per minute)
        reading_time_minutes = word_count / 200
        
        # Most common words (excluding common stop words)
        stop_words = {'the', 'a', 'an', 'and', 'or', 'but', 'in', 'on', 'at', 'to', 'for', 'of', 'with', 'by', 'is', 'are', 'was', 'were', 'be', 'been', 'have', 'has', 'had', 'do', 'does', 'did', 'will', 'would', 'could', 'should', 'this', 'that', 'these', 'those'}
        words = [word.lower().strip('.,!?";()[]{}') for word in text.split()]
        content_words = [word for word in words if word not in stop_words and len(word) > 2]
        word_freq = {}
        for word in content_words:
            word_freq[word] = word_freq.get(word, 0) + 1
        
        top_words = sorted(word_freq.items(), key=lambda x: x[1], reverse=True)[:5]
        
        analysis = f"""Text Analysis Results:
        
📊 Basic Statistics:
• Words: {word_count}
• Characters: {char_count} (with spaces), {char_count_no_spaces} (without spaces)
• Sentences: {sentence_count}
• Paragraphs: {paragraph_count}

📈 Averages:
• Words per sentence: {avg_words_per_sentence:.1f}
• Characters per word: {avg_chars_per_word:.1f}

⏱️ Reading Time: {reading_time_minutes:.1f} minutes

🔤 Most Common Words:
{chr(10).join([f"• {word}: {count}" for word, count in top_words])}
        """
        
        return analysis
        
    except Exception as e:
        return f"Text analysis error: {str(e)}"

@tool
def url_shortener(url: str, custom_alias: str = "") -> str:
    """Create a shortened URL using TinyURL service.
    
    Args:
        url: The URL to shorten
        custom_alias: Optional custom alias for the shortened URL
    """
    try:
        if not url.startswith(('http://', 'https://')):
            url = 'https://' + url
        
        api_url = "http://tinyurl.com/api-create.php"
        params = {'url': url}
        
        if custom_alias:
            params['alias'] = custom_alias
        
        response = requests.get(api_url, params=params, timeout=10)
        
        if response.status_code == 200:
            short_url = response.text.strip()
            if short_url.startswith('http'):
                return f"Shortened URL: {short_url}"
            else:
                return f"Error: {short_url}"
        else:
            return f"URL shortening failed with status code: {response.status_code}"
            
    except Exception as e:
        return f"URL shortening error: {str(e)}"

@tool
def password_generator(length: int = 12, include_symbols: bool = True, include_numbers: bool = True) -> str:
    """Generate a secure random password.
    
    Args:
        length: Length of the password (default: 12)
        include_symbols: Whether to include special characters (default: True)
        include_numbers: Whether to include numbers (default: True)
    """
    try:
        import random
        import string
        
        if length < 4:
            return "Error: Password length must be at least 4 characters"
        
        if length > 128:
            return "Error: Password length cannot exceed 128 characters"
        
        # Base character set
        chars = string.ascii_letters
        
        if include_numbers:
            chars += string.digits
        
        if include_symbols:
            chars += "!@#$%^&*()_+-=[]{}|;:,.<>?"
        
        # Ensure at least one character from each selected category
        password = []
        
        # Add at least one lowercase and uppercase letter
        password.append(random.choice(string.ascii_lowercase))
        password.append(random.choice(string.ascii_uppercase))
        
        if include_numbers:
            password.append(random.choice(string.digits))
        
        if include_symbols:
            password.append(random.choice("!@#$%^&*()_+-=[]{}|;:,.<>?"))
        
        # Fill the rest with random characters
        for _ in range(length - len(password)):
            password.append(random.choice(chars))
        
        # Shuffle the password
        random.shuffle(password)
        
        generated_password = ''.join(password)
        
        return f"Generated password: {generated_password}\n\nSecurity tips:\n• Store in a password manager\n• Don't reuse across sites\n• Change regularly for sensitive accounts"
        
    except Exception as e:
        return f"Password generation error: {str(e)}"

@tool
def unit_converter(value: float, from_unit: str, to_unit: str) -> str:
    """Convert between different units of measurement.
    
    Args:
        value: The numeric value to convert
        from_unit: Source unit (e.g., 'km', 'miles', 'kg', 'lbs', 'celsius', 'fahrenheit')
        to_unit: Target unit
    """
    try:
        # Conversion factors to base units
        conversions = {
            # Length (to meters)
            'mm': 0.001, 'cm': 0.01, 'm': 1, 'km': 1000,
            'inch': 0.0254, 'ft': 0.3048, 'yard': 0.9144, 'mile': 1609.34,
            
            # Weight (to kg)
            'mg': 0.000001, 'g': 0.001, 'kg': 1,
            'oz': 0.0283495, 'lb': 0.453592, 'stone': 6.35029,
            
            # Temperature (special handling)
            'celsius': 'celsius', 'fahrenheit': 'fahrenheit', 'kelvin': 'kelvin',
            
            # Volume (to liters)
            'ml': 0.001, 'l': 1, 'gallon': 3.78541, 'quart': 0.946353,
            'pint': 0.473176, 'cup': 0.236588, 'fl_oz': 0.0295735
        }
        
        from_unit = from_unit.lower()
        to_unit = to_unit.lower()
        
        # Handle temperature conversions separately
        if from_unit in ['celsius', 'fahrenheit', 'kelvin'] or to_unit in ['celsius', 'fahrenheit', 'kelvin']:
            if from_unit == 'celsius' and to_unit == 'fahrenheit':
                result = (value * 9/5) + 32
            elif from_unit == 'fahrenheit' and to_unit == 'celsius':
                result = (value - 32) * 5/9
            elif from_unit == 'celsius' and to_unit == 'kelvin':
                result = value + 273.15
            elif from_unit == 'kelvin' and to_unit == 'celsius':
                result = value - 273.15
            elif from_unit == 'fahrenheit' and to_unit == 'kelvin':
                result = (value - 32) * 5/9 + 273.15
            elif from_unit == 'kelvin' and to_unit == 'fahrenheit':
                result = (value - 273.15) * 9/5 + 32
            else:
                return f"Error: Cannot convert from {from_unit} to {to_unit}"
            
            return f"{value}° {from_unit.title()} = {result:.2f}° {to_unit.title()}"
        
        # Handle other unit conversions
        if from_unit not in conversions or to_unit not in conversions:
            available_units = list(conversions.keys())
            return f"Error: Unsupported unit. Available units: {', '.join(available_units)}"
        
        # Convert to base unit, then to target unit
        base_value = value * conversions[from_unit]
        result = base_value / conversions[to_unit]
        
        return f"{value} {from_unit} = {result:.6f} {to_unit}"
        
    except Exception as e:
        return f"Unit conversion error: {str(e)}"

@tool
def json_formatter(json_string: str) -> str:
    """Format and validate JSON strings.
    
    Args:
        json_string: The JSON string to format
    """
    try:
        # Parse the JSON to validate it
        parsed = json.loads(json_string)
        
        # Format with proper indentation
        formatted = json.dumps(parsed, indent=2, ensure_ascii=False)
        
        return f"✅ Valid JSON - Formatted:\n\n{formatted}"
        
    except json.JSONDecodeError as e:
        return f"❌ Invalid JSON - Error: {str(e)}"
    except Exception as e:
        return f"JSON formatting error: {str(e)}"

@tool
def base64_encoder_decoder(text: str, operation: str = "encode") -> str:
    """Encode or decode base64 strings.
    
    Args:
        text: The text to encode/decode
        operation: Either 'encode' or 'decode'
    """
    try:
        import base64
        
        if operation.lower() == "encode":
            encoded = base64.b64encode(text.encode('utf-8')).decode('utf-8')
            return f"Base64 Encoded: {encoded}"
        
        elif operation.lower() == "decode":
            decoded = base64.b64decode(text.encode('utf-8')).decode('utf-8')
            return f"Base64 Decoded: {decoded}"
        
        else:
            return "Error: Operation must be 'encode' or 'decode'"
            
    except Exception as e:
        return f"Base64 operation error: {str(e)}"

@tool
def hash_generator(text: str, algorithm: str = "sha256") -> str:
    """Generate hash for text using various algorithms.
    
    Args:
        text: The text to hash
        algorithm: Hash algorithm ('md5', 'sha1', 'sha256', 'sha512')
    """
    try:
        import hashlib
        
        algorithm = algorithm.lower()
        
        if algorithm == "md5":
            hash_obj = hashlib.md5(text.encode())
        elif algorithm == "sha1":
            hash_obj = hashlib.sha1(text.encode())
        elif algorithm == "sha256":
            hash_obj = hashlib.sha256(text.encode())
        elif algorithm == "sha512":
            hash_obj = hashlib.sha512(text.encode())
        else:
            return "Error: Supported algorithms are 'md5', 'sha1', 'sha256', 'sha512'"
        
        hash_value = hash_obj.hexdigest()
        return f"{algorithm.upper()} hash: {hash_value}"
        
    except Exception as e:
        return f"Hash generation error: {str(e)}"

# =============================================================================
# MAIN APPLICATION SETUP
# =============================================================================

def main():
    """Initialize and launch the AI agent application."""
    
    # Initialize core tools
    final_answer = FinalAnswerTool()
    search_tool = DuckDuckGoSearchTool()
    
    # Initialize the language model
    model = InferenceClientModel(
        max_tokens=2096,
        temperature=0.5,
        model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
        custom_role_conversions=None,
    )
    
    # Try to load image generation tool
    try:
        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"⚠️ Could not load image generation tool: {e}")
        image_generation_tool = None
    
    # Collect all tools
    tools = [
        final_answer,
        search_tool,
        get_current_time_in_timezone,
        weather_forecast,
        calculate_advanced_math,
        text_analyzer,
        url_shortener,
        password_generator,
        unit_converter,
        json_formatter,
        base64_encoder_decoder,
        hash_generator,
    ]
    
    # Add image generation tool if available
    if image_generation_tool:
        tools.append(image_generation_tool)
    
    # Load system prompt templates
    try:
        with open("prompts.yaml", 'r') as stream:
            prompt_templates = yaml.safe_load(stream)
        print("✅ System prompts loaded successfully")
    except FileNotFoundError:
        print("⚠️ prompts.yaml not found, using default prompts")
        prompt_templates = None
    except Exception as e:
        print(f"⚠️ Error loading prompts.yaml: {e}")
        prompt_templates = None
    
    # Initialize the CodeAgent
    agent = CodeAgent(
        model=model,
        tools=tools,
        max_steps=10,  # Increased for more complex tasks
        verbosity_level=1,
        grammar=None,
        planning_interval=3,  # Plan every 3 steps
        name="SuperAgent",
        description="An advanced AI agent capable of solving complex tasks using code and various tools",
        prompt_templates=prompt_templates
    )
    
    print("🚀 SuperAgent initialized with the following capabilities:")
    print("   • Web search and information retrieval")
    print("   • Advanced mathematical calculations")
    print("   • Text analysis and processing")
    print("   • Time zone and weather information")
    print("   • URL shortening and web utilities")
    print("   • Security tools (password generation, hashing)")
    print("   • Data conversion and formatting")
    if image_generation_tool:
        print("   • Image generation from text")
    print("   • Code execution and problem solving")
    print()
    
    # Launch the Gradio interface
    try:
        ui = GradioUI(agent)
        ui.launch(
            share=False,  # Set to True to create a public link
            inbrowser=True,  # Open in browser automatically
            server_name="0.0.0.0",  # Allow access from other devices on network
            server_port=7860,  # Default Gradio port
        )
    except Exception as e:
        print(f"❌ Failed to launch Gradio UI: {e}")
        print("Make sure you have gradio installed and Gradio_UI module is available")

if __name__ == "__main__":
    main()