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

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

        Project: {project_name}
        Description: {project_description}
        Use Cases: {actors_use_cases}

        Instructions:
        1. Show the main workflows as activities
        2. Include decision points and parallel activities where appropriate
        3. Use standard PlantUML activity diagram syntax
        4. Don't include any explanation, just the PlantUML code

        Example format (DO NOT TREAT THIS AS VARIABLES):
        '''plantuml
        @startuml
        start
        :Main Activity;
        if (Decision?) then (yes)
          :Activity 1;
        else (no)
          :Activity 2;
        endif
        stop
        @enduml
        '''
        """

        async for chunk in self._stream_process(
            state=state,
            prompt_template=prompt_template,
            output_key="activity_diagram",
            step_name="generate_activity_diagram",
            project_name=state["project_name"],
            project_description=state["project_description"],
            actors_use_cases=state["actors_use_cases"]
        ):
            yield chunk