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 CalculatorState(BaseModel):
    input: CalculatorInput
    output: CalculatorOutput = None

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")

class SearchState(BaseModel):
    input: SearchInput
    output: SearchOutput = None

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

    # Create the graph with state schema
    workflow = StateGraph(state_schema=CalculatorState)
    print("Calculator graph for workflow created")
    # Add the calculator tool node
    workflow.add_node("calculator", ToolNode(calculator_function))
    print("Calculator tool node added")
    # Set the entry and exit points
    workflow.set_entry_point("calculator")
    workflow.set_finish_point("calculator")
    print("Calculator workflow set")
    return workflow.compile()

def create_search_tool() -> Graph:
    """Creates a search tool using DuckDuckGo that can search for information online."""
    
    def search_function(state: SearchState) -> dict:
        with DDGS() as ddgs:
            # Run search
            raw_results = list(ddgs.text(
                state.input.query,
                max_results=state.input.max_results
            ))

        results = []
        for r in raw_results:
            try:
                results.append(SearchResult(
                    title=r.get("title", ""),
                    link=r.get("href", r.get("link", "")),  # DuckDuckGo sometimes uses "href"
                    snippet=r.get("body", r.get("snippet", ""))
                ))
            except Exception as e:
                print("Skipping malformed search result:", r, "Error:", e)

        return {
            "output": SearchOutput(
                results=results,
                query=state.input.query
            )
        }

    # Create the graph with state schema
    workflow = StateGraph(state_schema=SearchState)
    
    # 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}")