import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Try to import real smolagents, fallback to mock
try:
    from smolagents import ToolCallingAgent, tool
    from smolagents.models import LiteLLMModel
    SMOLAGENTS_AVAILABLE = True
    print("✅ smolagents framework loaded successfully!")
except ImportError:
    print("⚠️ smolagents not found. Using mock implementation for demonstration.")
    print("   To install: pip install git+https://github.com/huggingface/smolagents.git")
    from mock_smolagents import ToolCallingAgent, tool, LiteLLMModel
    SMOLAGENTS_AVAILABLE = False

# Define tools using the decorator (works with both real and mock)
@tool
def calculator(expression: str) -> str:
    """Evaluate a mathematical expression.
    
    Args:
        expression: A mathematical expression (e.g., '2 + 2 * 3')
    
    Returns:
        The result of the evaluation
    """
    try:
        # Safe evaluation
        allowed_chars = set('0123456789+-*/(). ')
        if not all(c in allowed_chars for c in expression):
            return "Error: Only basic math operations allowed: numbers, +, -, *, /, (, ), ."
        
        # Use eval with limited scope
        result = eval(expression, {"__builtins__": {}}, {})
        return f"Calculator: {expression} = {result}"
    except Exception as e:
        return f"Calculation error: {str(e)}"

@tool
def search_web(query: str) -> str:
    """Search the web for information.
    
    Args:
        query: Search query
    
    Returns:
        Search results summary
    """
    try:
        # Try to use duckduckgo if available
        from duckduckgo_search import DDGS
        
        results = []
        with DDGS() as ddgs:
            for r in ddgs.text(query, max_results=3):
                results.append(f"• {r['title']}: {r['body'][:100]}...")
        
        if results:
            return f"Web search for '{query}':\n" + "\n".join(results)
        return f"No web results found for '{query}'"
    except ImportError:
        return f"[Mock] Web search for '{query}':\n• Result 1: Information about {query}\n• Result 2: More details..."
    except Exception as e:
        return f"Search error: {str(e)}"

@tool
def get_hf_dataset_info(dataset_name: str = "mnist") -> str:
    """Get information about a Hugging Face dataset.
    
    Args:
        dataset_name: Name of the dataset (e.g., 'mnist', 'glue')
    
    Returns:
        Dataset information
    """
    try:
        from huggingface_hub import get_dataset_info
        
        info = get_dataset_info(dataset_name)
        response = f"📊 **Hugging Face Dataset: {info.id}**\n"
        response += f"📥 Downloads: {info.downloads:,}\n"
        response += f"📝 Description: {info.description[:250]}..."
        
        return response
    except ImportError:
        datasets = {
            "mnist": "MNIST: 70,000 handwritten digits (28x28 grayscale images), standard computer vision dataset",
            "glue": "GLUE: General Language Understanding Evaluation benchmark for natural language understanding"
        }
        info = datasets.get(dataset_name.lower(), f"Dataset '{dataset_name}' not found in mock database")
        return f"📊 **Dataset Info (Mock):**\n{info}"
    except Exception as e:
        return f"Error getting dataset info: {str(e)}"

@tool
def translate_text(text: str, target_language: str = "Spanish") -> str:
    """Translate text to another language.
    
    Args:
        text: Text to translate
        target_language: Target language (default: Spanish)
    
    Returns:
        Translated text
    """
    translations = {
        "hello": {
            "spanish": "hola",
            "french": "bonjour",
            "german": "hallo",
            "italian": "ciao"
        },
        "world": {
            "spanish": "mundo",
            "french": "monde",
            "german": "welt",
            "italian": "mondo"
        },
        "goodbye": {
            "spanish": "adiós",
            "french": "au revoir",
            "german": "auf wiedersehen",
            "italian": "addio"
        }
    }
    
    text_lower = text.lower()
    lang_lower = target_language.lower()
    
    if text_lower in translations and lang_lower in translations[text_lower]:
        return f"Translation: '{text}' → {target_language}: '{translations[text_lower][lang_lower]}'"
    else:
        return f"No translation found for '{text}' to {target_language}. Try: 'hello', 'world', or 'goodbye'"

class MyHuggingFaceAgent:
    def __init__(self, use_mock_model=False):
        """Initialize the agent.
        
        Args:
            use_mock_model: Use mock model instead of real API (for testing)
        """
        self.use_mock_model = use_mock_model or not SMOLAGENTS_AVAILABLE
        
        # Get all tools
        self.tools = [calculator, search_web, get_hf_dataset_info, translate_text]
        
        # Initialize the agent
        self.agent = self._create_agent()
        
        print(f"Agent initialized with {len(self.tools)} tools")
        print(f"smolagents available: {SMOLAGENTS_AVAILABLE}")
    
    def _create_agent(self):
        """Create the smolagents agent."""
        try:
            if self.use_mock_model or not SMOLAGENTS_AVAILABLE:
                # Use mock model
                model = LiteLLMModel(
                    model_id="mock-model",
                    api_key="mock-key"
                )
            else:
                # Use real model (requires API key)
                api_key = os.getenv("OPENAI_API_KEY") or os.getenv("HF_TOKEN")
                if api_key and api_key != "not-provided":
                    model = LiteLLMModel(
                        model_id="gpt-3.5-turbo",
                        api_key=api_key
                    )
                else:
                    print("⚠️ No API key found. Using mock model.")
                    model = LiteLLMModel(
                        model_id="mock-model",
                        api_key="mock-key"
                    )
            
            # Create agent
            agent = ToolCallingAgent(
                model=model,
                tools=self.tools,
                max_steps=5,
                verbose=True,
                add_base_tools=False
            )
            
            return agent
            
        except Exception as e:
            print(f"Error creating agent: {e}")
            # Return a basic agent as fallback
            return None
    
    def run(self, task: str) -> str:
        """Run the agent on a task."""
        try:
            if self.agent:
                # Use smolagents agent
                result = self.agent.run(task)
                
                # Format the response
                response = "🤖 **AI Agent Response (smolagents Framework)**\n\n"
                response += f"{result}\n\n"
                response += "---\n"
                response += "🛠️ **Framework:** smolagents (Hugging Face)\n"
                response += f"📊 **Status:** {'Using real framework' if SMOLAGENTS_AVAILABLE else 'Using mock for demonstration'}"
                
                return response
            else:
                # Fallback to simple processing
                return self._simple_agent(task)
                
        except Exception as e:
            return f"⚠️ **Agent Error:** {str(e)}\n\n{self._simple_agent(task)}"
    
    def _simple_agent(self, task: str) -> str:
        """Simple fallback agent."""
        response = "🔧 **Simple Agent Response**\n\n"
        
        # Try each tool
        for tool_func in self.tools:
            tool_name = tool_func.name
            if tool_name in task.lower():
                try:
                    # Simple argument extraction
                    if tool_name == "calculator":
                        import re
                        expr = re.sub(r'[^\d\+\-\*\/\(\)\.\s]', '', task).strip()
                        if expr:
                            result = tool_func(expr)
                        else:
                            result = "Please provide a math expression"
                    elif tool_name == "search_web":
                        query = task.replace("search", "").replace("find", "").strip()
                        result = tool_func(query if query else task)
                    elif tool_name == "get_hf_dataset_info":
                        result = tool_func()
                    elif tool_name == "translate_text":
                        if " to " in task.lower():
                            parts = task.lower().split(" to ")
                            text = parts[0].replace("translate", "").strip()
                            lang = parts[1].strip().title()
                            result = tool_func(text, lang)
                        else:
                            result = tool_func("hello")
                    else:
                        result = tool_func(task)
                    
                    response += f"**{tool_name.replace('_', ' ').title()}:**\n{result}\n\n"
                    break
                except Exception as e:
                    response += f"Error using {tool_name}: {str(e)}\n\n"
        else:
            # No tool matched
            response += """**I can help with:**
1. **Calculator** - Perform math calculations
   Example: 'calculate 15 * 3' or 'what is 45 + 23?'
2. **Web Search** - Search for information
   Example: 'search for AI news' or 'find machine learning tutorials'
3. **Dataset Info** - Get Hugging Face dataset information
   Example: 'tell me about mnist dataset' or 'glue dataset info'
4. **Translation** - Translate words
   Example: 'translate hello to Spanish' or 'translate goodbye to French'
Try one of these commands!"""
        
        return response