from enum import Enum

from langchain_core.messages.tool import tool_call
from langchain_core.tools import tool
from langgraph.constants import START, END
from langgraph.graph import StateGraph
from pydantic.dataclasses import dataclass
from uipath.tracing import traced


class Operator(Enum):
    ADD = "+"
    SUBTRACT = "-"
    MULTIPLY = "*"
    DIVIDE = "/"

@dataclass
class CalculatorInput:
    a: float
    b: float
    operator: Operator

@dataclass
class CalculatorOutput:
    result: float

@tool
async def example_tool(x: int) -> int:
    """
    Mock tool
    """
    return x


@traced(name="postprocess")
async def postprocess(x: float) -> float:
    """
    Example of nested traced invocation.
    """
    return x

@traced(name="calculate")
async def calculate(input: CalculatorInput) -> CalculatorOutput:
    result = 0
    match input.operator:
        case Operator.ADD: result = input.a + input.b
        case Operator.SUBTRACT: result = input.a - input.b
        case Operator.MULTIPLY: result = input.a * input.b
        case Operator.DIVIDE: result = input.a / input.b if input.b != 0 else 0
    result = await postprocess(result)
    result = await example_tool.ainvoke(input={"x": result})
    return CalculatorOutput(result=result)

builder = StateGraph(state_schema=CalculatorInput, input=CalculatorInput, output=CalculatorOutput)

builder.add_node("calculate", calculate)
builder.add_edge(START, "calculate")
builder.add_edge("calculate", END)

graph = builder.compile()

@tool
@traced(name="example_tool")
async def ...