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

# --- Calculator Tool Types ---

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

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

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

# --- Search Tool Types ---

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
    output: Optional[SearchOutput]

# --- Calculator Tool Implementation ---

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"]
            }
        }

    workflow = StateGraph(state_schema=CalculatorState)
    workflow.add_node("calculator", ToolNode(calculator_function))
    workflow.set_entry_point("calculator")
    workflow.set_finish_point("calculator")
    print("Calculator workflow created and compiled")
    return workflow.compile()

# --- Search Tool Implementation ---

def create_search_tool() -> Graph:
    """Creates a search tool using DuckDuckGo that can search for information online."""
    print("Creating search tool")

    def search_function(state: SearchState) -> Dict[str, Any]:
        print("Search function called")
        query = state["input"]["query"]
        max_results = state["input"].get("max_results", 3)

        with DDGS() as ddgs:
            raw_results = list(ddgs.text(query, max_results=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": query
            }
        }

    workflow = StateGraph(state_schema=SearchState)
    workflow.add_node("search", ToolNode(search_function))
    workflow.set_entry_point("search")
    workflow.set_finish_point("search")
    print("Search workflow created and compiled")
    return workflow.compile()