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

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

        Project: {project_name}
        Description: {project_description}

        Instructions:
        1. Show the physical deployment architecture
        2. Include nodes, artifacts, and their relationships
        3. Use standard PlantUML deployment diagram syntax
        4. Don't include any explanation, just the PlantUML code

        Example format (DO NOT TREAT THIS AS VARIABLES):
        '''plantuml
        @startuml
        node "Server" as server {{
          artifact "WebApp"
        }}
        node "Client" as client
        client --> server : HTTP
        @enduml
        '''
        """

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