from typing import Dict, Any, List
from langgraph.graph import Graph, StateGraph
from langgraph.prebuilt import ToolNode
from pydantic import BaseModel, Field
from duckduckgo_search import DDGS

class CalculatorInput(BaseModel):
    operation: str = Field(..., description="The operation to perform (add, subtract, multiply, divide)")
    numbers: List[float] = Field(..., description="List of numbers to perform the operation on")

class CalculatorOutput(BaseModel):
    result: float = Field(..., description="The result of the calculation")
    operation: str = Field(..., description="The operation that was performed")

class SearchInput(BaseModel):
    query: str = Field(..., description="The search query to look up")
    max_results: int = Field(default=3, description="Maximum number of results to return")

class SearchResult(BaseModel):
    title: str = Field(..., description="Title of the search result")
    link: str = Field(..., description="URL of the search result")
    snippet: str = Field(..., description="Brief description of the search result")

class SearchOutput(BaseModel):
    results: List[SearchResult] = Field(..., description="List of search results")
    query: str = Field(..., description="The original search query")

def create_calculator_tool() -> Graph:
    """Creates a calculator tool using LangGraph that can perform basic arithmetic operations."""
    
    def calculator_function(state: Dict[str, Any]) -> Dict[str, Any]:
        input_data = CalculatorInput(**state["input"])
        
        if len(input_data.numbers) < 2:
            raise ValueError("At least two numbers are required for calculation")
        
        result = input_data.numbers[0]
        
        for num in input_data.numbers[1:]:
            if input_data.operation == "add":
                result += num
            elif input_data.operation == "subtract":
                result -= num
            elif input_data.operation == "multiply":
                result *= num
            elif input_data.operation == "divide":
                if num == 0:
                    raise ValueError("Cannot divide by zero")
                result /= num
            else:
                raise ValueError(f"Unsupported operation: {input_data.operation}")
        
        output = CalculatorOutput(
            result=result,
            operation=input_data.operation
        )
        
        return {"output": output}

    # Create the graph
    workflow = StateGraph()
    
    # Add the calculator tool node
    workflow.add_node("calculator", ToolNode(calculator_function))
    
    # Set the entry and exit points
    workflow.set_entry_point("calculator")
    workflow.set_finish_point("calculator")
    
    return workflow.compile()

def create_search_tool() -> Graph:
    """Creates a search tool using DuckDuckGo that can search for information online."""
    
    def search_function(state: Dict[str, Any]) -> Dict[str, Any]:
        input_data = SearchInput(**state["input"])
        
        # Initialize DuckDuckGo search
        with DDGS() as ddgs:
            # Perform the search
            search_results = list(ddgs.text(
                input_data.query,
                max_results=input_data.max_results
            ))
            
            # Convert results to our model
            results = [
                SearchResult(
                    title=result["title"],
                    link=result["link"],
                    snippet=result["body"]
                )
                for result in search_results
            ]
            
            output = SearchOutput(
                results=results,
                query=input_data.query
            )
            
            return {"output": output}

    # Create the graph
    workflow = StateGraph()
    
    # Add the search tool node
    workflow.add_node("search", ToolNode(search_function))
    
    # Set the entry and exit points
    workflow.set_entry_point("search")
    workflow.set_finish_point("search")
    
    return workflow.compile()

# Example usage:
# if __name__ == "__main__":
#     # Create the calculator tool
#     calculator = create_calculator_tool()
    
#     # Example calculation
#     result = calculator.invoke({
#         "input": {
#             "operation": "add",
#             "numbers": [1, 2, 3, 4]
#         }
#     })
    
#     print(f"Result: {result['output'].result}")