"""
Code Review A2A Agent - Built with a2a-python SDK (NOT Synqed)

This is a REAL A2A agent built using the a2a-python SDK.
It demonstrates that Synqed can route to agents built with ANY framework,
as long as they implement the A2A protocol.

This agent:
- Built with a2a-python SDK (not Synqed)
- Exposes A2A endpoints (AgentCard, message/send, etc.)
- Can be hosted anywhere
- Synqed routes to it via RemoteA2AAgent

Run: python code_review_a2a_agent.py
Then it will be available at http://localhost:8001
"""

import asyncio
import os
import logging
from dotenv import load_dotenv

# A2A SDK imports
from a2a.server.apps import A2AStarletteApplication
from a2a.server.request_handlers import DefaultRequestHandler
from a2a.server.tasks import InMemoryTaskStore
from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.events import EventQueue
from a2a.utils import new_agent_text_message
from a2a.types import (
    AgentCapabilities,
    AgentCard,
    AgentSkill,
)

load_dotenv()
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class CodeReviewAgent:
    """
    Core agent logic for code review.
    
    This agent uses an LLM to perform code review analysis.
    Built with standard Python + LLM calls (not Synqed).
    """
    
    def __init__(self):
        """Initialize the code review agent."""
        self.name = "CodeReviewAgent"
    
    async def review(self, message: str) -> str:
        """
        Perform code review on the given message/request.
        
        Args:
            message: The user's request or code to review
            
        Returns:
            Code review analysis as a string
        """
        # Use Anthropic Claude for code review
        import anthropic
        
        client = anthropic.AsyncAnthropic(
            api_key=os.getenv("ANTHROPIC_API_KEY")
        )
        
        prompt = f"""You are a code review expert agent.

User request: {message}

Provide insightful code review feedback. If the user provides code, analyze it.
If they describe a concept, explain code quality principles related to it.

Keep your response concise (2-3 paragraphs).

Return JSON format: {{"send_to": "Coordinator", "content": "Your analysis"}}"""
        
        response = await client.messages.create(
            model="claude-sonnet-4-5",
            max_tokens=500,
            messages=[{"role": "user", "content": prompt}]
        )
        
        return response.content[0].text.strip()


class CodeReviewAgentExecutor(AgentExecutor):
    """
    A2A Agent Executor for Code Review.
    
    This implements the A2A AgentExecutor interface, making our agent
    compatible with the A2A protocol.
    """
    
    def __init__(self):
        """Initialize the executor with the core agent."""
        self.agent = CodeReviewAgent()
        logger.info("CodeReviewAgentExecutor initialized")
    
    async def execute(
        self,
        context: RequestContext,
        event_queue: EventQueue,
    ) -> None:
        """
        Execute the agent logic.
        
        This is called when the agent receives a message via A2A protocol.
        
        Args:
            context: Request context containing the incoming message
            event_queue: Queue for emitting response events
        """
        logger.info("Executing code review agent")
        
        # Extract the user message from context
        # context.message contains the A2A Message object
        user_message = ""
        if context.message and context.message.parts:
            for part in context.message.parts:
                if hasattr(part, 'text') and part.text:
                    user_message += part.text
        
        logger.info(f"Received message: {user_message[:100]}...")
        
        # Process with the agent
        result = await self.agent.review(user_message)
        
        logger.info(f"Generated response: {result[:100]}...")
        
        # Emit the response via event queue
        # This will be picked up by the A2A server and sent back to the caller
        await event_queue.enqueue_event(new_agent_text_message(result))
    
    async def cancel(
        self,
        context: RequestContext,
        event_queue: EventQueue
    ) -> None:
        """
        Cancel task execution.
        
        For this simple agent, we don't support cancellation.
        """
        logger.warning("Cancel requested but not supported")
        raise Exception("Cancel not supported for CodeReviewAgent")


def create_agent_card(host: str, port: int) -> AgentCard:
    """
    Create the AgentCard for this agent.
    
    The AgentCard is published at /.well-known/agent-card.json
    and describes the agent's capabilities, skills, and how to call it.
    
    Args:
        host: Host where agent is running
        port: Port where agent is running
        
    Returns:
        AgentCard describing this agent
    """
    skill = AgentSkill(
        id='code_review',
        name='Code Review',
        description='Analyzes code and provides quality feedback',
        tags=['code', 'review', 'quality', 'analysis'],
        examples=[
            'Review this Python function',
            'Analyze code quality',
            'What are best practices for error handling?'
        ],
    )
    
    return AgentCard(
        name='Code Review Agent',
        description='Expert code review agent built with A2A SDK (not Synqed)',
        url=f'http://{host}:{port}/',
        version='1.0.0',
        default_input_modes=['text'],
        default_output_modes=['text'],
        capabilities=AgentCapabilities(
            streaming=True,
            push_notifications=False,
        ),
        skills=[skill],
    )


def main(host: str = 'localhost', port: int = 8001):
    """
    Start the Code Review A2A Agent server.
    
    This creates an A2A-compliant HTTP server that:
    - Publishes AgentCard at /.well-known/agent-card.json
    - Accepts messages at POST / (JSON-RPC)
    - Returns structured A2A responses
    
    Args:
        host: Host to bind to
        port: Port to bind to
    """
    logger.info(f"Starting Code Review A2A Agent on {host}:{port}")
    
    # Create agent card
    agent_card = create_agent_card(host, port)
    logger.info(f"Agent card created: {agent_card.name}")
    
    # Create request handler with our agent executor
    request_handler = DefaultRequestHandler(
        agent_executor=CodeReviewAgentExecutor(),
        task_store=InMemoryTaskStore(),
    )
    logger.info("Request handler created")
    
    # Create A2A Starlette application
    server = A2AStarletteApplication(
        agent_card=agent_card,
        http_handler=request_handler,
    )
    logger.info("A2A server application created")
    
    # Build and run the server
    app = server.build()
    
    logger.info(f"🚀 Code Review A2A Agent running at http://{host}:{port}")
    logger.info(f"📄 AgentCard available at http://{host}:{port}/.well-known/agent-card.json")
    logger.info("✅ Ready to receive A2A messages from Synqed or any A2A client")
    
    import uvicorn
    uvicorn.run(app, host=host, port=port, log_level="info")


if __name__ == '__main__':
    main()

