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

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

        Project: {project_name}
        Description: {project_description}
        Sequence: {sequence_interactions}

        Instructions:
        1. Show the timing constraints between components
        2. Include lifelines and state changes over time
        3. Use standard PlantUML timing diagram syntax
        4. Don't include any explanation, just the PlantUML code

        Example format (DO NOT TREAT THIS AS VARIABLES):
        '''plantuml
        @startuml
        clock "Clock" as C with period 1
        binary "Signal" as S
        @0
        S is low
        @5
        S is high
        @10
        S is low
        @enduml
        '''
        """

        async for chunk in self._stream_process(
            state=state,
            prompt_template=prompt_template,
            output_key="timing_diagram",
            step_name="generate_timing_diagram",
            project_name=state["project_name"],
            project_description=state["project_description"],
            sequence_interactions=state["sequence_interactions"]
        ):
            yield chunk