from typing import Annotated

from pydantic import BaseModel, ValidationError

from agentdojo.functions_runtime import Depends, FunctionsRuntime, TaskEnvironment

# --8<-- [start:stateless_function]
runtime = FunctionsRuntime()


# Define a function
@runtime.register_function  # (1)!
def add(a: int, b: int) -> int:  # (2)!
    """Sums two numbers.

    :param a: The first number.
    :param b: The second number.
    """  # (3)!
    return a + b


# --8<-- [end:stateless_function]

# --8<-- [start:run_stateless_function]

result, _ = runtime.run_function(
    None,  # (1)!
    "add",
    {"a": 1, "b": 2},
)
assert result == 3
# --8<-- [end:run_stateless_function]


# --8<-- [start:error]
_, error = runtime.run_function(
    None,
    "add",
    {"a": 1, "b": "hello"},
)
assert error is not None
assert "ValidationError" in error  # (1)!
# --8<-- [end:error]


# --8<-- [start:exception]
try:
    _, error = runtime.run_function(
        None,
        "add",
        {"a": 1, "b": "hello"},
        raise_on_error=True,
    )
except ValidationError:
    ...
else:
    assert False, "Expected a validation error"

# --8<-- [end:exception]


# --8<-- [start:stateful_function]
class Counter(BaseModel):
    value: int

    def increment(self) -> int:
        self.value += 1
        return self.value

    def increment_by(self, n: int) -> int:
        self.value += n
        return self.value


class Environment(TaskEnvironment):
    counter: Counter


@runtime.register_function
def increment_counter(
    counter: Annotated[Counter, Depends("counter")],  # (1)!
) -> int:
    """Increments the counter by 1."""  # (2)!
    return counter.increment()


def increment_counter_by(
    counter: Annotated[Counter, Depends("counter")],
    n: int,
) -> int:
    """Increments the counter by `n`.

    :param n: The number to increment the counter by.
    """
    return counter.increment_by(n)


runtime.register_function(increment_counter_by)  # (3)!

# --8<-- [end:stateful_function]


# --8<-- [start:run_stateful_function]

environment = Environment(counter=Counter(value=1))
result, _ = runtime.run_function(
    environment,  # (1)!
    "increment_counter_by",
    {"n": 5},
)
assert result == 6
# --8<-- [end:run_stateful_function]
