"""
Author: Vasiliy Zdanovskiy
email: vasilyvz@gmail.com

Registration management utilities for MCP Proxy Adapter API.
"""

from __future__ import annotations

import asyncio
from typing import Any, Dict, List, Optional

from mcp_proxy_adapter.client.jsonrpc_client import JsonRpcClient
from mcp_proxy_adapter.core.logging import get_global_logger

from mcp_proxy_adapter.api.core.registration_context import (
    HeartbeatSettings,
    ProxyCredentials,
    RegistrationContext,
    prepare_registration_context,
    resolve_heartbeat_settings,
    resolve_runtime_credentials,
    resolve_unregister_endpoint,
)
from mcp_proxy_adapter.api.core.registration_tasks import (
    create_heartbeat_task,
    unregister_from_proxy as unregister_task,
)

# Global registration status with mutex for thread-safe access
_registration_status_lock = asyncio.Lock()
_registration_status: bool = False


async def get_registration_status() -> bool:
    """Get current registration status (thread-safe with mutex)."""
    async with _registration_status_lock:
        return _registration_status


async def set_registration_status(status: bool) -> None:
    """Set registration status (thread-safe with mutex)."""
    async with _registration_status_lock:
        global _registration_status
        _registration_status = status


class RegistrationManager:
    """Manager for proxy registration functionality using JsonRpcClient."""

    def __init__(self) -> None:
        """Initialize registration manager."""
        self.logger = get_global_logger()
        self.registered = False
        self._registration_config: Optional[Dict[str, Any]] = None
        self.registration_task: Optional[asyncio.Task] = None
        self.server_name: Optional[str] = None
        self.server_url: Optional[str] = None
        self.proxy_url: Optional[str] = None
        self.capabilities: List[str] = []
        self.metadata: Dict[str, Any] = {}
        self.config: Optional[Dict[str, Any]] = None
        self._use_proxy_client = False
        self._proxy_client_config: Dict[str, Any] = {}
        self._proxy_registration_config: Dict[str, Any] = {}
        self._registration_credentials: Optional[ProxyCredentials] = None
        self._runtime_credentials: Optional[ProxyCredentials] = None
        self._register_endpoint: str = "/register"
        self._heartbeat_settings: Optional[HeartbeatSettings] = None

    async def register_with_proxy(self, config: Dict[str, Any]) -> bool:
        """
        Register this server with the proxy using JsonRpcClient.

        Supports both ``proxy_client`` (SimpleConfig format) and ``proxy_registration`` (legacy format).
        Registration is controlled by ``registration.auto_on_startup`` rather than ``proxy_client.enabled``.
        """

        context = prepare_registration_context(config, self.logger)
        if context is None:
            return True

        self._apply_context(context, config)

        proxy_url = self.proxy_url
        assert proxy_url is not None
        assert self.server_name is not None
        assert self.server_url is not None

        # Parse proxy URL to extract protocol, host, and port
        from urllib.parse import urlparse

        parsed = urlparse(proxy_url)
        proxy_protocol = parsed.scheme or "http"
        proxy_host = parsed.hostname or "localhost"
        proxy_port = parsed.port or (443 if proxy_protocol == "https" else 80)

        # Extract client certificate paths from credentials if available
        client_cert = None
        client_key = None
        client_ca = None
        if context.credentials.cert:
            client_cert, client_key = context.credentials.cert
        if isinstance(context.credentials.verify, str):
            client_ca = context.credentials.verify

        self.logger.debug(
            "Creating JsonRpcClient: protocol=%s, host=%s, port=%s, cert=%s, ca=%s",
            proxy_protocol,
            proxy_host,
            proxy_port,
            client_cert is not None,
            client_ca is not None,
        )
        if client_cert:
            self.logger.debug("   Client cert: %s, key: %s", client_cert, client_key)
        if client_ca:
            self.logger.debug("   CA cert: %s", client_ca)

        client = JsonRpcClient(
            protocol=proxy_protocol,
            host=proxy_host,
            port=proxy_port,
            cert=client_cert,
            key=client_key,
            ca=client_ca,
        )

        async def _register() -> Dict[str, Any]:
            self._log_credentials("🔐 Registration SSL config", context.credentials)
            self.logger.info(f"📡 Connecting to proxy: {proxy_url}")
            self.logger.debug(
                "   Endpoint: %s, Server: %s -> %s",
                self._register_endpoint,
                self.server_name,
                self.server_url,
            )
            # Note: cert and verify are already set in JsonRpcClient constructor
            # but we pass them again for explicit override if needed
            return await client.register_with_proxy(
                proxy_url=proxy_url,
                server_name=context.server_name,
                server_url=context.advertised_url,
                capabilities=self.capabilities,
                metadata=self.metadata,
                cert=context.credentials.cert,
                verify=context.credentials.verify,
            )

        max_retries = 5
        retry_delay = 2

        try:
            for attempt in range(max_retries):
                try:
                    registration_response = await _register()
                    if registration_response is not None:
                        self.logger.debug(
                            "Proxy registration response payload: %s",
                            registration_response,
                        )
                        self.logger.info(
                            "✅ Successfully registered with proxy as %s -> %s",
                            self.server_name,
                            self.server_url,
                        )
                        self.registered = True
                        # Update global registration status with mutex
                        await set_registration_status(True)
                        return True
                except Exception as exc:  # noqa: BLE001
                    full_error = self._format_httpx_error(exc)
                    if attempt < max_retries - 1:
                        self.logger.warning(
                            "⚠️  Registration attempt %s/%s failed: %s. Retrying in %ss...",
                            attempt + 1,
                            max_retries,
                            full_error,
                            retry_delay,
                        )
                        await asyncio.sleep(retry_delay)
                        retry_delay *= 2
                    else:
                        self.logger.error(
                            "❌ Failed to register with proxy after %s attempts: %s",
                            max_retries,
                            full_error,
                        )
                        return False
            # If we get here, all retries failed
            return False
        except Exception as exc:  # noqa: BLE001
            self.logger.error(f"❌ Registration error: {exc}")
            return False
        finally:
            await client.close()

    async def start_heartbeat(self, config: Dict[str, Any]) -> None:
        """Start heartbeat task using JsonRpcClient.

        Heartbeat task will handle registration if not yet registered,
        then send heartbeats periodically.
        """

        # Store config for registration in heartbeat loop
        self._registration_config = config

        # Prepare context to get proxy URL and settings
        # Note: context may be None if auto_on_startup=false, but we still want to start heartbeat
        # for potential future registration or just heartbeat if already registered
        context = prepare_registration_context(config, self.logger)
        self.logger.info(f"🔍 Registration context prepared: {context is not None}")
        
        # If context is None, try to prepare minimal context for heartbeat only
        if context is None:
            self.logger.info("⚠️  Registration context is None, creating minimal context for heartbeat")
            # Check if proxy_client is enabled (heartbeat should still run even if auto_on_startup=false)
            proxy_client_config = dict(config.get("proxy_client") or {})
            if not proxy_client_config.get("enabled", False):
                self.logger.debug(
                    "Proxy client not enabled, heartbeat task will not start"
                )
                return

            # Try to create minimal context for heartbeat (without registration requirement)
            # This allows heartbeat to run even if auto_on_startup=false
            # Note: resolve_heartbeat_settings and resolve_runtime_credentials are already imported at module level

            # Extract minimal info for heartbeat
            self._use_proxy_client = True
            self._proxy_client_config = proxy_client_config
            self._proxy_registration_config = {}

            # Set minimal server info for heartbeat
            server_config = dict(config.get("server") or {})
            host = server_config.get("host", "127.0.0.1")
            port = server_config.get("port", 8000)
            protocol = server_config.get("protocol", "http")
            advertised_host = server_config.get("advertised_host") or host
            scheme = "https" if protocol in ("https", "mtls") else "http"
            self.server_url = f"{scheme}://{advertised_host}:{port}"
            self.server_name = (
                proxy_client_config.get("server_id")
                or proxy_client_config.get("server_name")
                or f"mcp-adapter-{host}-{port}"
            )

            # Build proxy URL
            pc_protocol = proxy_client_config.get("protocol", "http")
            pc_scheme = "https" if pc_protocol in ("https", "mtls") else "http"
            pc_host = proxy_client_config.get("host", "localhost")
            pc_port = proxy_client_config.get("port", 3005)
            self.proxy_url = f"{pc_scheme}://{pc_host}:{pc_port}"

            self.capabilities = proxy_client_config.get(
                "capabilities", ["jsonrpc", "health"]
            )
            self.metadata = {}
        else:
            # Normal path: context is available
            self.logger.info("✅ Registration context available, applying context...")
            self._apply_context(context, config)
            self.logger.info(f"✅ Context applied: server_name={self.server_name}, proxy_url={self.proxy_url}")
            
            # Check after applying context
            if not self._can_start_tasks():
                self.logger.warning(
                    "⚠️  Cannot start tasks (missing proxy_url, server_name, or server_url)"
                )
                return

        credentials = resolve_runtime_credentials(
            self._use_proxy_client,
            self._proxy_client_config,
            self._proxy_registration_config,
        )
        settings = resolve_heartbeat_settings(
            self._use_proxy_client,
            self._proxy_client_config,
            self._proxy_registration_config,
        )
        self._runtime_credentials = credentials
        self._heartbeat_settings = settings

        heartbeat_url = f"{self.proxy_url}{settings.endpoint}"
        self.logger.info(
            "💓 Starting heartbeat task (interval: %ss)", settings.interval
        )

        assert self.server_name is not None
        assert self.server_url is not None

        self.registration_task = create_heartbeat_task(
            registration_manager=self,
            proxy_url=heartbeat_url,
            server_name=self.server_name,
            server_url=self.server_url,
            capabilities=self.capabilities,
            metadata=self.metadata,
            settings=settings,
            credentials=credentials,
            logger=self.logger,
        )

    async def stop(self) -> None:
        """Stop registration manager and unregister from proxy."""

        if self.registration_task:
            self.registration_task.cancel()
            try:
                await self.registration_task
            except asyncio.CancelledError:
                pass
            self.registration_task = None

        if not (self.registered and self._can_start_tasks() and self.config):
            self.registered = False
            return

        credentials = self._runtime_credentials or resolve_runtime_credentials(
            self._use_proxy_client,
            self._proxy_client_config,
            self._proxy_registration_config,
        )
        endpoint = resolve_unregister_endpoint(
            self._use_proxy_client,
            self._proxy_client_config,
        )

        assert self.proxy_url is not None
        assert self.server_name is not None

        try:
            await unregister_task(
                proxy_url=self.proxy_url,
                server_name=self.server_name,
                endpoint=endpoint,
                credentials=credentials,
                logger=self.logger,
            )
        except Exception as exc:  # noqa: BLE001
            self.logger.error(f"Error unregistering from proxy: {exc}")
        finally:
            self.registered = False

    def _apply_context(
        self, context: RegistrationContext, config: Dict[str, Any]
    ) -> None:
        self.server_name = context.server_name
        self.server_url = context.advertised_url
        self.proxy_url = context.proxy_url
        self.capabilities = list(context.capabilities)
        self.metadata = dict(context.metadata)
        self.config = config
        self._use_proxy_client = context.use_proxy_client
        self._proxy_client_config = context.proxy_client_config
        self._proxy_registration_config = context.proxy_registration_config
        self._registration_credentials = context.credentials
        self._runtime_credentials = None
        self._register_endpoint = context.register_endpoint

    def _log_credentials(self, prefix: str, credentials: ProxyCredentials) -> None:
        self.logger.info(
            "%s: cert=%s, verify=%s",
            prefix,
            credentials.cert is not None,
            credentials.verify,
        )
        if credentials.cert:
            self.logger.debug("   Client cert: %s, key: %s", *credentials.cert)
        if isinstance(credentials.verify, str):
            self.logger.debug("   CA cert: %s", credentials.verify)

    def _format_httpx_error(self, exc: Exception) -> str:
        import httpx

        error_msg = str(exc) or type(exc).__name__
        details: List[str] = [f"type={type(exc).__name__}"]

        if isinstance(exc, httpx.HTTPStatusError):
            details.append(f"status={exc.response.status_code}")
            try:
                details.append(f"response={exc.response.text[:200]}")
            except Exception:  # noqa: BLE001
                pass
        elif isinstance(exc, httpx.ConnectError):
            details.append("connection_failed")
            if hasattr(exc, "request"):
                details.append(f"url={exc.request.url}")
        elif isinstance(exc, httpx.TimeoutException):
            details.append("timeout")

        return f"{error_msg} ({', '.join(details)})" if details else error_msg

    def _can_start_tasks(self) -> bool:
        return bool(self.proxy_url and self.server_name and self.server_url)
