from typing import Dict, Any, List, TypedDict, Optional
from langgraph.graph import Graph, StateGraph
from langgraph.prebuilt import ToolNode
from duckduckgo_search import DDGS

class CalculatorInput(TypedDict):
    operation: str
    numbers: List[float]

class CalculatorOutput(TypedDict):
    result: float
    operation: str

class CalculatorState(TypedDict):
    input: CalculatorInput
    output: Optional[CalculatorOutput]

class SearchInput(TypedDict):
    query: str
    max_results: int

class SearchResult(TypedDict):
    title: str
    link: str
    snippet: str

class SearchOutput(TypedDict):
    results: List[SearchResult]
    query: str

class SearchState(TypedDict):
    input: SearchInput  # Now typed correctly
    output: Optional[SearchOutput]

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[str, Any]:
        print("Calculator function called")
        input_data = 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']}")
        
        return {
            "output": {
                "result": result,
                "operation": input_data["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[str, Any]:
        with DDGS() as ddgs:
            raw_results = list(ddgs.text(
                state["input"]["query"],
                max_results=state["input"]["max_results"]
            ))

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

        return {
            "output": {
                "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}")