"""Main MCP Proxy orchestration.

This module provides the main McpProxy class that orchestrates all components:
- Stdio interface for communication with Kiro
- Protocol translator for stdio <-> HTTP conversion
- Authenticated HTTP client for backend communication
- Configuration management
- Graceful startup and shutdown
- Signal handling and cleanup

Requirements satisfied:
- 3.1: Perform OAuth DCR with OAuth provider on startup
- 2.1: Establish connections to HTTP MCP server
"""

import asyncio
import logging
import signal
from typing import Optional

from .auth.manager import AuthenticationManagerImpl
from .config.manager import ConfigurationManagerImpl
from .http.authenticated_client import AuthenticatedHttpClient
from .http.client import HttpClientImpl
from .models import JsonRpcMessage, ProxyConfig
from .translator.translator import ProtocolTranslatorImpl
from .stdio.interface import StdioInterface
from .exceptions import McpProxyError, ConfigurationError
from .logging_config import get_logger, configure_logging

logger = get_logger(__name__)


class McpProxy:
    """Main MCP Proxy orchestrator.
    
    This class brings together all components to create a complete proxy that:
    - Presents a stdio MCP interface to Kiro
    - Translates between stdio and HTTP MCP protocols
    - Manages OAuth DCR authentication
    - Handles HTTP communication with backend MCP server
    - Provides graceful startup and shutdown
    
    The proxy operates as follows:
    1. Load configuration from environment variables
    2. Initialize authentication manager (performs DCR)
    3. Create HTTP client with authentication
    4. Set up protocol translator
    5. Start stdio interface
    6. Process messages bidirectionally
    7. Handle shutdown signals gracefully
    
    Attributes:
        config: Proxy configuration
        is_running: Whether the proxy is currently running
    """

    def __init__(self, config: Optional[ProxyConfig] = None):
        """Initialize the MCP Proxy.
        
        Args:
            config: Optional proxy configuration. If not provided,
                   configuration will be loaded from environment variables.
        """
        self._config = config
        self._running = False
        self._shutdown_event = asyncio.Event()
        
        # Components (initialized during startup)
        self._config_manager: Optional[ConfigurationManagerImpl] = None
        self._auth_manager: Optional[AuthenticationManagerImpl] = None
        self._http_client: Optional[HttpClientImpl] = None
        self._authenticated_client: Optional[AuthenticatedHttpClient] = None
        self._translator: Optional[ProtocolTranslatorImpl] = None
        self._stdio_interface: Optional[StdioInterface] = None
        
        logger.info("McpProxy instance created")

    @property
    def config(self) -> Optional[ProxyConfig]:
        """Get the current configuration."""
        return self._config

    @property
    def is_running(self) -> bool:
        """Check if the proxy is currently running."""
        return self._running

    async def start(self) -> None:
        """Start the MCP Proxy.
        
        This method performs the complete startup sequence:
        1. Load configuration (if not provided)
        2. Set up logging
        3. Initialize authentication manager (performs DCR)
        4. Create and initialize HTTP client
        5. Set up protocol translator
        6. Start stdio interface
        7. Install signal handlers
        
        Raises:
            ConfigurationError: If configuration is invalid
            AuthenticationError: If OAuth DCR fails
            RuntimeError: If proxy is already running
        """
        if self._running:
            raise RuntimeError("McpProxy is already running")
        
        logger.info("Starting MCP Proxy")
        
        try:
            # Step 1: Load configuration
            await self._load_configuration()
            
            # Step 2: Set up logging
            self._setup_logging()
            
            # Step 3: Initialize authentication manager
            await self._initialize_authentication()
            
            # Step 4: Create and initialize HTTP client
            await self._initialize_http_client()
            
            # Step 5: Set up protocol translator
            self._initialize_translator()
            
            # Step 6: Start stdio interface
            await self._start_stdio_interface()
            
            # Step 7: Install signal handlers
            self._install_signal_handlers()
            
            self._running = True
            self._shutdown_event.clear()
            
            logger.info("MCP Proxy started successfully")
            
        except Exception as e:
            logger.error(f"Failed to start MCP Proxy: {e}", exc_info=True)
            # Clean up any partially initialized components
            await self._cleanup()
            raise

    async def stop(self) -> None:
        """Stop the MCP Proxy gracefully.
        
        This method performs graceful shutdown:
        1. Signal shutdown to all components
        2. Stop stdio interface (drains pending messages)
        3. Close authenticated HTTP client
        4. Close authentication manager
        5. Clean up resources
        """
        if not self._running:
            logger.warning("McpProxy is not running")
            return
        
        logger.info("Stopping MCP Proxy")
        self._running = False
        self._shutdown_event.set()
        
        try:
            # Stop stdio interface first (drains pending messages)
            if self._stdio_interface:
                await self._stdio_interface.stop()
            
            # Close authenticated HTTP client
            if self._authenticated_client:
                await self._authenticated_client.close()
            
            # Close authentication manager
            if self._auth_manager:
                await self._auth_manager.close()
            
            # Clean up remaining resources
            await self._cleanup()
            
            logger.info("MCP Proxy stopped successfully")
            
        except Exception as e:
            logger.error(f"Error during shutdown: {e}", exc_info=True)
            raise

    async def run(self) -> None:
        """Run the proxy until shutdown signal is received.
        
        This method starts the proxy and waits for a shutdown signal.
        It's the main entry point for running the proxy as a standalone process.
        
        Usage:
            proxy = McpProxy()
            await proxy.run()
        """
        await self.start()
        
        try:
            # Wait for shutdown signal
            logger.info("MCP Proxy running, waiting for shutdown signal")
            await self._shutdown_event.wait()
        except KeyboardInterrupt:
            logger.info("Received keyboard interrupt")
        finally:
            await self.stop()

    async def wait_for_shutdown(self) -> None:
        """Wait for shutdown signal.
        
        This method blocks until the proxy is signaled to shut down.
        Useful for keeping the main process alive while the proxy runs.
        """
        await self._shutdown_event.wait()

    # ========================================================================
    # Private Methods - Initialization
    # ========================================================================

    async def _load_configuration(self) -> None:
        """Load configuration from environment variables if not provided."""
        if self._config is None:
            logger.info("Loading configuration from environment variables")
            self._config_manager = ConfigurationManagerImpl()
            self._config = await self._config_manager.load()
            logger.info("Configuration loaded successfully")
        else:
            logger.info("Using provided configuration")
        
        # Log configuration (excluding sensitive data)
        logger.info(
            "Configuration: mcp_server_url=%s, oauth_provider_url=%s, "
            "client_name=%s, scopes=%s, connection_timeout=%s, "
            "retry_attempts=%s, log_level=%s",
            self._config.mcp_server_url,
            self._config.oauth_provider_url,
            self._config.client_name,
            self._config.scopes,
            self._config.connection_timeout,
            self._config.retry_attempts,
            self._config.log_level,
        )

    def _setup_logging(self) -> None:
        """Set up logging based on configuration."""
        if self._config:
            configure_logging(self._config.log_level)
            logger.info(f"Logging configured with level: {self._config.log_level}")

    async def _initialize_authentication(self) -> None:
        """Initialize authentication manager and perform DCR.
        
        This satisfies Requirement 3.1: Perform OAuth DCR on startup.
        """
        if not self._config:
            raise ConfigurationError("Configuration not loaded")
        
        logger.info("Initializing authentication manager")
        self._auth_manager = AuthenticationManagerImpl(self._config)
        
        # Initialize performs DCR and obtains initial token
        await self._auth_manager.initialize()
        
        logger.info("Authentication manager initialized successfully")

    async def _initialize_http_client(self) -> None:
        """Initialize HTTP client with authentication.
        
        This satisfies Requirement 2.1: Establish connections to HTTP MCP server.
        """
        if not self._config or not self._auth_manager:
            raise RuntimeError("Configuration or authentication not initialized")
        
        logger.info("Initializing HTTP client")
        
        # Create retry configuration
        from .http.client import RetryConfig
        retry_config = RetryConfig(
            max_retries=self._config.retry_attempts,
            max_delay=self._config.max_backoff_seconds,
        )
        
        # Create base HTTP client
        self._http_client = HttpClientImpl(
            base_url=str(self._config.mcp_server_url),
            timeout=self._config.connection_timeout,
            retry_config=retry_config,
        )
        
        # Wrap with authentication
        self._authenticated_client = AuthenticatedHttpClient(
            http_client=self._http_client,
            auth_manager=self._auth_manager,
        )
        
        # Initialize authenticated client (ensures token is ready)
        await self._authenticated_client.initialize()
        
        logger.info("HTTP client initialized successfully")

    def _initialize_translator(self) -> None:
        """Initialize protocol translator."""
        if not self._config:
            raise RuntimeError("Configuration not initialized")
        
        logger.info("Initializing protocol translator")
        
        self._translator = ProtocolTranslatorImpl(
            base_url=str(self._config.mcp_server_url)
        )
        
        logger.info("Protocol translator initialized successfully")

    async def _start_stdio_interface(self) -> None:
        """Start stdio interface with message handler."""
        logger.info("Starting stdio interface")
        
        # Create stdio interface with message handler
        self._stdio_interface = StdioInterface(
            message_handler=self._handle_stdio_message
        )
        
        # Start the interface
        await self._stdio_interface.start()
        
        # Wait for interface to be fully started
        started = await self._stdio_interface.wait_until_started(timeout=5.0)
        if not started:
            raise RuntimeError("Stdio interface failed to start within timeout")
        
        logger.info("Stdio interface started successfully")

    # ========================================================================
    # Private Methods - Message Handling
    # ========================================================================

    async def _handle_stdio_message(
        self,
        message: JsonRpcMessage
    ) -> Optional[JsonRpcMessage]:
        """Handle incoming stdio message from Kiro.
        
        This is the main message processing pipeline:
        1. Translate stdio message to HTTP request
        2. Send HTTP request via authenticated client
        3. Translate HTTP response back to stdio message
        4. Return stdio response
        
        Args:
            message: JSON-RPC message from Kiro
            
        Returns:
            JSON-RPC response message, or None for notifications
        """
        if not self._translator or not self._authenticated_client:
            logger.error(
                "Components not initialized",
                has_translator=self._translator is not None,
                has_authenticated_client=self._authenticated_client is not None,
            )
            return self._create_error_response(
                message.id,
                -32603,
                "Internal error: proxy not fully initialized"
            )
        
        try:
            # Handle notifications (no response expected)
            if message.is_notification():
                logger.debug(
                    "Received notification",
                    method=message.method,
                )
                # For now, we don't forward notifications to backend
                # This could be extended to support server-to-client notifications
                return None
            
            # Translate stdio request to HTTP
            logger.debug(
                "Translating stdio request to HTTP",
                method=message.method,
                id=message.id,
                has_params=message.params is not None,
            )
            http_request = await self._translator.translate_stdio_to_http(message)
            
            # Send HTTP request via authenticated client
            logger.debug(
                "Sending HTTP request",
                http_method=http_request.method,
                url=http_request.url,
                session_id=http_request.session_id,
            )
            http_response = await self._authenticated_client.send_request(http_request)
            
            # Translate HTTP response back to stdio
            logger.debug(
                "Translating HTTP response to stdio",
                status=http_response.status,
                content_type=http_response.content_type,
                body_length=len(http_response.body) if http_response.body else 0,
            )
            stdio_response = await self._translator.translate_http_to_stdio(
                http_response,
                stdio_request_id=message.id
            )
            
            # Mark correlation as completed
            if message.id is not None:
                try:
                    self._translator.mark_correlation_completed(message.id)
                except Exception as e:
                    logger.warning(
                        "Failed to mark correlation completed",
                        error=str(e),
                        message_id=message.id,
                    )
            
            logger.info(
                "Message processed successfully",
                method=message.method,
                request_id=message.id,
                response_has_result=stdio_response.result is not None,
                response_has_error=stdio_response.error is not None,
            )
            return stdio_response
            
        except McpProxyError as e:
            logger.error(
                "Proxy error handling message",
                error=str(e),
                error_type=type(e).__name__,
                message_id=message.id,
                method=message.method,
                has_details=hasattr(e, 'details') and e.details is not None,
                exc_info=True,
            )
            
            # Mark correlation as failed
            if message.id is not None:
                try:
                    self._translator.mark_correlation_failed(message.id)
                except Exception:
                    pass
            
            # Return error response
            from .exceptions import to_jsonrpc_error
            error_obj = to_jsonrpc_error(e)
            return self._create_error_response(
                message.id,
                error_obj["code"],
                error_obj["message"],
                error_obj.get("data")
            )
            
        except Exception as e:
            logger.error(
                "Unexpected error handling message",
                error=str(e),
                error_type=type(e).__name__,
                message_id=message.id,
                method=message.method,
                exc_info=True,
            )
            
            # Mark correlation as failed
            if message.id is not None:
                try:
                    self._translator.mark_correlation_failed(message.id)
                except Exception:
                    pass
            
            # Return generic error response
            return self._create_error_response(
                message.id,
                -32603,
                f"Internal error: {str(e)}"
            )

    def _create_error_response(
        self,
        request_id: Optional[int | str],
        code: int,
        message: str,
        data: Optional[any] = None
    ) -> JsonRpcMessage:
        """Create a JSON-RPC error response.
        
        Args:
            request_id: Request ID to respond to
            code: JSON-RPC error code
            message: Error message
            data: Optional additional error data
            
        Returns:
            JSON-RPC error response
        """
        from .models import JsonRpcError
        
        return JsonRpcMessage(
            jsonrpc="2.0",
            id=request_id,
            error=JsonRpcError(
                code=code,
                message=message,
                data=data
            )
        )

    # ========================================================================
    # Private Methods - Signal Handling
    # ========================================================================

    def _install_signal_handlers(self) -> None:
        """Install signal handlers for graceful shutdown."""
        loop = asyncio.get_event_loop()
        
        for sig in (signal.SIGTERM, signal.SIGINT):
            try:
                loop.add_signal_handler(
                    sig,
                    lambda s=sig: asyncio.create_task(self._handle_signal(s))
                )
                logger.debug(f"Installed signal handler for {sig.name}")
            except NotImplementedError:
                # Signal handlers not supported on this platform (e.g., Windows)
                logger.warning(f"Signal handlers not supported for {sig.name}")

    async def _handle_signal(self, sig: signal.Signals) -> None:
        """Handle shutdown signals.
        
        Args:
            sig: Signal that was received
        """
        logger.info(f"Received signal {sig.name}, initiating shutdown")
        await self.stop()

    # ========================================================================
    # Private Methods - Cleanup
    # ========================================================================

    async def _cleanup(self) -> None:
        """Clean up resources during shutdown or after failed startup."""
        logger.debug("Cleaning up resources")
        
        # Close components in reverse order of initialization
        if self._stdio_interface and self._stdio_interface.is_running:
            try:
                await self._stdio_interface.stop()
            except Exception as e:
                logger.error(f"Error stopping stdio interface: {e}")
        
        if self._authenticated_client:
            try:
                await self._authenticated_client.close()
            except Exception as e:
                logger.error(f"Error closing authenticated client: {e}")
        
        if self._auth_manager:
            try:
                await self._auth_manager.close()
            except Exception as e:
                logger.error(f"Error closing auth manager: {e}")
        
        # Clear references
        self._stdio_interface = None
        self._authenticated_client = None
        self._http_client = None
        self._auth_manager = None
        self._translator = None
        
        logger.debug("Cleanup completed")

    # ========================================================================
    # Context Manager Support
    # ========================================================================

    async def __aenter__(self) -> "McpProxy":
        """Async context manager entry."""
        await self.start()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """Async context manager exit."""
        await self.stop()

    def __repr__(self) -> str:
        """String representation of the proxy."""
        return f"McpProxy(running={self._running}, config={self._config is not None})"
