"""
Server wrapper for simplified agent server creation and management.
"""

import asyncio
import logging
from typing import Any

from a2a.server.apps.jsonrpc import A2AFastAPIApplication
from a2a.server.request_handlers import DefaultRequestHandler
from a2a.server.tasks import InMemoryTaskStore
from a2a.types import AgentCard

from synqed.agent import Agent

logger = logging.getLogger(__name__)


# Filter to suppress CancelledError logging during server shutdown
class _CancelledErrorFilter(logging.Filter):
    """Suppresses CancelledError tracebacks that occur during graceful shutdown."""
    
    def filter(self, record: logging.LogRecord) -> bool:
        if record.levelno == logging.ERROR:
            # Check exception info
            if record.exc_info:
                exc_type = record.exc_info[0]
                if exc_type and exc_type.__name__ == "CancelledError":
                    return False
            # Also check if message contains traceback with CancelledError
            msg = str(record.getMessage())
            if "CancelledError" in msg or "asyncio.exceptions.CancelledError" in msg:
                return False
        return True


# Apply the filter to all relevant loggers to prevent noisy shutdown errors
_error_filter = _CancelledErrorFilter()
logging.getLogger("uvicorn.error").addFilter(_error_filter)
logging.getLogger("uvicorn").addFilter(_error_filter)

# Also reduce logging level for the root uvicorn error logger during shutdown scenarios
# This is a fallback to catch cases where the filter doesn't work
def _suppress_cancellation_errors():
    """Configure logging to suppress CancelledError messages."""
    # Get all uvicorn-related loggers
    for logger_name in ["uvicorn.error", "uvicorn"]:
        log = logging.getLogger(logger_name)
        # Add our filter to all handlers
        for handler in log.handlers:
            handler.addFilter(_error_filter)

_suppress_cancellation_errors()


class AgentServer:
    """
    Simplified server wrapper for hosting A2A agents.
    
    This class provides an easy way to create and start a server for your agent
    without dealing with the underlying A2A server configuration.
    
    Example:
        ```python
        agent = Agent(
            name="My Agent",
            description="Does something",
            skills=["skill1"]
        )
        
        server = AgentServer(agent, port=8000)
        await server.start()  # Blocking
        
        # Or run in background
        await server.start_background()
        # ... do other things ...
        await server.stop()
        ```
    """
    
    def __init__(
        self,
        agent: Agent,
        host: str = "0.0.0.0",
        port: int = 8000,
        path_prefix: str = "",
        task_store: Any | None = None,
        enable_cors: bool = True,
    ):
        """
        Initialize the agent server.
        
        Args:
            agent: The Agent instance to serve
            host: Host to bind to (default: 0.0.0.0)
            port: Port to bind to (default: 8000)
            path_prefix: Path prefix for endpoints (default: empty string)
            task_store: Optional custom task store (default: InMemoryTaskStore)
            enable_cors: Whether to enable CORS (default: True)
        """
        self.agent = agent
        self.host = host
        self.port = port
        self.path_prefix = path_prefix
        
        # Update agent URL to match server configuration
        # Use localhost for client connections if binding to 0.0.0.0
        client_host = "localhost" if host == "0.0.0.0" else host
        self.agent.set_url(f"http://{client_host}:{port}{path_prefix}")
        
        # Create task store
        self.task_store = task_store or InMemoryTaskStore()
        
        # Create request handler
        self.request_handler = DefaultRequestHandler(
            agent_executor=self.agent.get_executor(),
            task_store=self.task_store,
        )
        
        # Create the A2A FastAPI application
        a2a_app = A2AFastAPIApplication(
            agent_card=self.agent.card,
            http_handler=self.request_handler,
        )
        
        # Build the FastAPI app
        self.app = a2a_app.build()
        
        # Add a welcome endpoint
        self._add_welcome_endpoint()
        
        # Enable CORS if requested
        if enable_cors:
            self._enable_cors()
        
        # Server process reference
        self._server_task: asyncio.Task | None = None
        self._server_started = False
        self._uvicorn_server: Any | None = None
    
    def _add_welcome_endpoint(self) -> None:
        """Add a welcome GET endpoint at the root."""
        @self.app.get("/")
        async def welcome():
            return {
                "agent": self.agent.name,
                "description": self.agent.description,
                "skills": self.agent.skills,
                "status": "online",
                "message": "This is an A2A agent server. Use POST requests to interact with the agent."
            }
        
        @self.app.get("/favicon.ico")
        async def favicon():
            """Handle favicon requests to prevent 404 errors."""
            from fastapi.responses import Response
            return Response(status_code=204)
    
    def _enable_cors(self) -> None:
        """Enable CORS for the FastAPI app."""
        try:
            from fastapi.middleware.cors import CORSMiddleware
            
            self.app.add_middleware(
                CORSMiddleware,
                allow_origins=["*"],
                allow_credentials=True,
                allow_methods=["*"],
                allow_headers=["*"],
            )
            logger.info("CORS enabled for agent server")
        except ImportError:
            logger.warning("FastAPI not available, CORS not enabled")
    
    async def start(self) -> None:
        """
        Start the server (blocking).
        
        This will run the server until interrupted.
        """
        try:
            import uvicorn
        except ImportError as e:
            raise ImportError(
                "uvicorn is required to run the server. "
                "Install with: pip install uvicorn"
            ) from e
        
        logger.info(
            f"Starting agent server for '{self.agent.name}' on "
            f"http://{self.host}:{self.port}{self.path_prefix}"
        )
        
        config = uvicorn.Config(
            app=self.app,
            host=self.host,
            port=self.port,
            log_level="info",
        )
        server = uvicorn.Server(config)
        self._server_started = True
        
        try:
            await server.serve()
        finally:
            self._server_started = False
    
    async def start_background(self) -> None:
        """
        Start the server in the background.
        
        Returns immediately while the server runs in a background task.
        Use stop() to shut down the server.
        """
        if self._server_task is not None:
            logger.warning("Server is already running in background")
            return
        
        self._server_task = asyncio.create_task(self._run_background())
        
        # Wait a bit for server to start
        await asyncio.sleep(1)
        
        logger.info(
            f"Agent server started in background for '{self.agent.name}' on "
            f"http://{self.host}:{self.port}{self.path_prefix}"
        )
    
    async def _run_background(self) -> None:
        """Internal method to run server in background."""
        try:
            import uvicorn
        except ImportError as e:
            raise ImportError(
                "uvicorn is required to run the server. "
                "Install with: pip install uvicorn"
            ) from e
        
        config = uvicorn.Config(
            app=self.app,
            host=self.host,
            port=self.port,
            log_level="info",
        )
        server = uvicorn.Server(config)
        self._uvicorn_server = server
        self._server_started = True
        
        try:
            await server.serve()
        finally:
            self._server_started = False
            self._server_task = None
            self._uvicorn_server = None
    
    async def stop(self) -> None:
        """Stop the background server."""
        if self._server_task is None:
            logger.warning("Server is not running")
            return
        
        # Shutdown uvicorn gracefully
        if self._uvicorn_server is not None:
            self._uvicorn_server.should_exit = True
            # Give it a moment to shutdown gracefully
            await asyncio.sleep(0.1)
        
        # Cancel the server task
        if self._server_task and not self._server_task.done():
            self._server_task.cancel()
            try:
                await self._server_task
            except asyncio.CancelledError:
                pass
        
        logger.info(f"Server stopped for agent '{self.agent.name}'")
        self._server_task = None
        self._server_started = False
        self._uvicorn_server = None
    
    @property
    def is_running(self) -> bool:
        """Check if the server is currently running."""
        return self._server_started
    
    @property
    def url(self) -> str:
        """Get the server URL."""
        return f"http://{self.host}:{self.port}{self.path_prefix}"
    
    def get_card(self) -> AgentCard:
        """Get the agent's card."""
        return self.agent.card
    
    def __repr__(self) -> str:
        """String representation of the server."""
        status = "running" if self.is_running else "stopped"
        return (
            f"AgentServer(agent='{self.agent.name}', "
            f"url='{self.url}', status='{status}')"
        )

