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

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

        Workflow Analysis: {sequence_interactions}

        Instructions:
        1. Use proper participant types:
           - actor for human users
           - boundary for UI elements  
           - control for controllers
           - entity for domain objects
           - database for storage
        2. Include all message types:
           - Synchronous: ->
           - Asynchronous: ->>
           - Return: -->
           - Self: -> self
        3. Add activation bars with:
           - activate/deactivate
           - Nested activations where needed
        4. Use advanced features:
           - Groups (alt/opt/loop/par)
           - Notes
           - Dividers (== Sections ==)
           - Stereotypes <<component>>
           - Arrows with:
             * Lost messages (->x)
             * Found messages (x->)
             * Parallel (->||)
        5. Format for readability:
           - Proper indentation
           - Section comments
           - Color coding where helpful

        Example Structure (DO NOT TREAT AS VARIABLES):
        @startuml
        skinparam style strictuml
        skinparam lifelineStrategy solid
        
        actor Customer <<Human>>
        boundary "Web UI" as UI
        control OrderController
        entity Order
        database OrderDB
        
        == Order Submission ==
        Customer -> UI: Submit Order
        activate UI
        UI -> OrderController: processOrder()
        activate OrderController
        
        group Payment Processing
            OrderController -> PaymentService <<External>>: authorizePayment()
            activate PaymentService
            PaymentService --> OrderController: status
            deactivate PaymentService
        end
        
        alt payment approved
            OrderController -> OrderDB: persistOrder()
            OrderController --> UI: confirmation
        else payment declined
            OrderController --> UI: error
        end
        
        deactivate OrderController
        UI --> Customer: Order Status
        deactivate UI
        @enduml
        """

        async for chunk in self._stream_process(
            state=state,
            prompt_template=prompt_template,
            output_key="sequence_diagram",
            step_name="generate_sequence_diagram",
            sequence_interactions=state["sequence_interactions"]
        ):
            yield chunk