"""This is an automatically generated file. Do not modify it.

This file was generated using `langgraph-gen` version {{ version }}.
To regenerate this file, run `langgraph-gen` with the source `yaml` file as an argument.

Usage:

1. Add the generated file to your project.
2. Create a new agent using the stub.

Below is a sample implementation of the generated stub:

```python
from typing_extensions import TypedDict

from [some_path] import {{ stub_name }}

class SomeState(TypedDict):
    # define your attributes here
    foo: str

# Define stand-alone functions
{% for node in nodes %}
def {{ node.id }}(state: SomeState) -> dict:
    print("In node: {{ node.name }}")
    return {
        # Add your state update logic here
    }


{% endfor %}
{% for edge in edges if edge.condition is defined %}
def {{ edge.condition }}(state: SomeState) -> str:
    print("In condition: {{ edge.condition }}")
    raise NotImplementedError("Implement me.")


{% endfor %}
agent = {{ stub_name }}(
    state_schema=SomeState,
    impl=[
    {% for node in nodes %}
        ("{{ node.name }}", {{ node.id }}),
    {% endfor %}
    {% for edge in edges if edge.condition is defined %}
        ("{{ edge.condition }}", {{ edge.condition }}),
    {% endfor %}
    ]
)

compiled_agent = agent.compile()

print(compiled_agent.invoke({"foo": "bar"}))
"""

from typing import Callable, Any, Optional, Type

from langgraph.constants import END
from langgraph.graph import StateGraph


def {{ stub_name }}(
    *,
    state_schema: Optional[Type[Any]] = None,
    config_schema: Optional[Type[Any]] = None,
    input: Optional[Type[Any]] = None,
    output: Optional[Type[Any]] = None,
    impl: list[tuple[str, Callable]],
) -> StateGraph:
    """Create the state graph for {{ stub_name }}."""
    # Declare the state graph
    builder = StateGraph(
        state_schema, config_schema=config_schema, input=input, output=output
    )

    nodes_by_name = {name: imp for name, imp in impl}

    all_names = set(nodes_by_name)

    expected_implementations = {
        {% for node in nodes %}
        "{{ node.id }}",
        {% endfor %}
        {% for edge in edges if edge.condition %}
        "{{ edge.condition }}",
        {% endfor %}
    }

    missing_nodes = expected_implementations - all_names
    if missing_nodes:
        raise ValueError(f"Missing implementations for: {missing_nodes}")

    extra_nodes = all_names - expected_implementations

    if extra_nodes:
        raise ValueError(
            f"Extra implementations for: {extra_nodes}. Please regenerate the stub."
        )

    # Add nodes
    {% for node in nodes %}
    builder.add_node("{{ node.name }}", nodes_by_name["{{ node.id }}"])
    {% endfor %}

    # Add edges
    {% for edge in edges %}
    {% if edge.condition is defined %}
    builder.add_conditional_edges(
        "{{ edge.from }}",
        nodes_by_name["{{ edge.condition }}"],
        {% if edge.paths is mapping %}
        {
        {% for key, value in edge.paths.items() %}
            "{{ key }}": {% if value == "__end__" %}END{% else %}"{{ value }}"{% endif %},
        {% endfor %}
        }
        {% else %}
        [
        {%  for path in edge.paths %}
            {% if value == "__end__" %}
            END,
            {% else %}
            "{{ path }}",
            {% endif %}
        {% endfor %}
        ],
        {% endif %}
    )
    {% else %}
    builder.add_edge("{{ edge.from }}", {% if edge.to == "__end__" %}END{% else %}"{{ edge.to }}"{% endif %})
    {% endif %}
    {% endfor %}

    builder.set_entry_point("{{ entrypoint }}")
    return builder

