from base_agent import BaseAgent
from models import AgentState
from typing import AsyncGenerator
import json

class StateDiagramGenerator(BaseAgent):
    async def __call__(self, state: AgentState) -> AsyncGenerator[str, None]:
        prompt_template = """
        Generate PlantUML code for a state diagram based on the following project:

        Project: {project_name}
        Description: {project_description}
        Entities: {entities_classes}

        Instructions:
        1. Show the state transitions for one key entity
        2. Include states, transitions, and events
        3. Use standard PlantUML state diagram syntax
        4. Don't include any explanation, just the PlantUML code

        Example format (DO NOT TREAT THIS AS VARIABLES):
        '''plantuml
        @startuml
        [*] --> State1
        State1 --> State2 : Event1
        State2 --> [*]
        @enduml
        '''
        """

        async for chunk in self._stream_process(
            state=state,
            prompt_template=prompt_template,
            output_key="state_diagram",
            step_name="generate_state_diagram",
            project_name=state["project_name"],
            project_description=state["project_description"],
            entities_classes=state["entities_classes"]
        ):
            yield chunk