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

class UseCaseDiagramGenerator(BaseAgent):
    async def __call__(self, state: AgentState) -> AsyncGenerator[str, None]:
        prompt_template = """
        Generate comprehensive PlantUML code for a use case diagram based on:
        
        Project: {project_name}
        Actors and Use Cases: {actors_use_cases}

        Instructions:
        1. Include all actors with proper typing:
           - Primary actors on left
           - Supporting actors on right
           - System actors if applicable
        2. Organize use cases into logical packages/subsystems when appropriate
        3. Use proper relationship types:
           - Association: --> (actor to use case)
           - Include: .> (base use case includes another)
           - Extend: .> (use case extends another with conditions)
        4. Apply stereotypes where helpful (<<system>>, <<external>>, etc.)
        5. Include notes for complex relationships when needed
        6. Use left-to-right direction for better readability
        7. Format with clear alignment and spacing
        8. Include only raw PlantUML code without explanations

        Example Structure (DO NOT TREAT AS VARIABLES):
        @startuml
        left to right direction
        skinparam packageStyle rectangle
        
        actor "Primary Actor" <<Human>> as user
        actor "Supporting System" <<System>> as legacy
        
        rectangle "Order Processing" {{
          usecase "Place Order" as UC1
          usecase "Process Payment" as UC2
          usecase "Validate Order" as UC3
          
          UC1 .> UC2 : include
          UC1 .> UC3 : include
        }}
        
        rectangle "Inventory" {{
          usecase "Check Stock" as UC4
        }}
        
        user --> UC1
        legacy --> UC3
        
        note right of UC1 
          This use case initiates the 
          order fulfillment workflow
        end note
        @enduml
        """

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