from typing import TypedDict
from src.config.llm import model
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.prebuilt import create_react_agent
from langgraph_swarm import create_handoff_tool, create_swarm
from .prompt import roleplay_prompt, guiding_prompt


class State(TypedDict):
    pass


def create_agents(scenario, checkpointer=InMemorySaver()):

    roleplay_agent = create_react_agent(
        model,
        [
            create_handoff_tool(
                agent_name="Guiding Agent",
                description="Hand off to Guiding Agent when user shows signs of needing help, guidance, or struggles with communication",
            ),
        ],
        prompt=roleplay_prompt.format(
            scenario_title=scenario["scenario_title"],
            scenario_description=scenario["scenario_description"],
            scenario_context=scenario["scenario_context"],
            your_role=scenario["your_role"],
            key_vocabulary=scenario["key_vocabulary"],
        ),
        name="Roleplay Agent",
    )

    guiding_agent = create_react_agent(
        model,
        [
            create_handoff_tool(
                agent_name="Roleplay Agent",
                description="Hand off back to Roleplay Agent when user is ready for scenario practice and shows improved confidence",
            ),
        ],
        prompt=guiding_prompt.format(
            scenario_title=scenario["scenario_title"],
            scenario_description=scenario["scenario_description"],
            scenario_context=scenario["scenario_context"],
            your_role=scenario["your_role"],
            key_vocabulary=scenario["key_vocabulary"],
        ),
        name="Guiding Agent",
    )

    workflow = create_swarm(
        [roleplay_agent, guiding_agent], default_active_agent="Roleplay Agent"
    )

    return workflow.compile(checkpointer)