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

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

        {entities_classes}

        Instructions:
        1. Include all classes with their complete attributes (visibility, type, constraints)
        2. Use proper UML relationships with correct cardinality and labels:
           - Inheritance: <|--
           - Implementation: <|..
           - Composition: *--
           - Aggregation: o--
           - Dependency: -->
        3. For complex relationships, use association classes when needed
        4. Include proper visibility markers:
           - Public: +
           - Private: -
           - Protected: #
           - Package: ~
        5. Use advanced PlantUML features where appropriate:
           - Abstract classes/interfaces
           - Enumerations
           - Stereotypes
           - Notes and constraints
           - Method signatures
        6. Organize classes into packages if the domain suggests logical groupings
        7. Use proper formatting with alignment and spacing
        8. Include only the raw PlantUML code without explanations

        Example Structure (DO NOT TREAT AS VARIABLES):
        @startuml
        package "Domain Model" {{
          abstract class AbstractEntity {{
            + id: UUID {{readonly}}
            # createdAt: DateTime
            + equals(other: Object): boolean {{abstract}}
          }}

          class User {{
            + username: String {{unique}}
            + email: String
            - passwordHash: String
            + authenticate(password: String): boolean
          }}

          enum UserRole {{
            ADMIN
            MEMBER
            GUEST
          }}

          User "1" *-- "0..*" Order : places >
          User o-- UserRole : has >

          (User, Order) .. OrderAssociation
          class OrderAssociation {{
            + orderDate: DateTime
            + status: OrderStatus
          }}
        }}

        note top of User
          User represents all system users
          with authentication capabilities
        end note
        @enduml
        """

        async for chunk in self._stream_process(
            state=state,
            prompt_template=prompt_template,
            output_key="class_diagram",
            step_name="generate_class_diagram",
            entities_classes=state["entities_classes"]
        ):
            yield chunk