"""
Session management for Daisi SDK.

This module provides SessionManager functionality similar to .NET's
SessionManagerBase, handling session negotiation, Orc redirection,
and Direct Connect support.

Key differences from .NET:
- Python uses composition over inheritance (no generics needed)
- Context managers for resource cleanup instead of IDisposable
- No factory pattern (Python's duck typing makes this unnecessary)
"""

import logging
from typing import Generic, Optional, TypeVar

import grpc

from protos.v1.models import HostModels_pb2, OrcModels_pb2, SessionModels_pb2

from .clients.session import SessionClient
from .config import DaisiConfig
from .exceptions import DaisiError

logger = logging.getLogger(__name__)


T = TypeVar("T")


class SessionManager:
    """Manages Daisi sessions with automatic negotiation and Direct Connect.

    The SessionManager handles:
    - Session creation and negotiation with the Orc
    - Automatic Orc redirection (MoveToOrc)
    - Direct Connect setup when the host supports it
    - Session lifecycle management

    Unlike .NET's SessionManagerBase<T> which uses generics for the client type,
    Python's SessionManager uses composition - you create the appropriate client
    after the session is established.

    Example:
        manager = SessionManager(config)
        manager.negotiate_session(host_id="my-host")

        if manager.use_direct_connect:
            # Create client using direct connection
            client = InferenceClient(config, direct_address=manager.direct_connect_address)
        else:
            # Use fully orchestrated connection
            client = InferenceClient(config, session_id=manager.session_id)
    """

    def __init__(
        self,
        config: Optional[DaisiConfig] = None,
        host_id: Optional[str] = None,
    ) -> None:
        """Initialize session manager.

        Args:
            config: Daisi configuration.
            host_id: Optional host ID to bind this manager to.
        """
        self.config = config or DaisiConfig()
        self.host_id = host_id
        self._session_id: Optional[str] = None
        self._host: Optional[HostModels_pb2.Host] = None
        self._connected = False
        self._use_direct_connect = False
        self._direct_connect_address: Optional[str] = None
        self._direct_connect_port: Optional[int] = None
        self._session_client: Optional[SessionClient] = None

    @property
    def session_id(self) -> Optional[str]:
        """Get the current session ID."""
        return self._session_id

    @property
    def host(self) -> Optional[HostModels_pb2.Host]:
        """Get the connected host."""
        return self._host

    @property
    def use_direct_connect(self) -> bool:
        """Check if Direct Connect is active."""
        return self._use_direct_connect

    @property
    def direct_connect_address(self) -> Optional[str]:
        """Get the Direct Connect host address."""
        return self._direct_connect_address

    @property
    def direct_connect_port(self) -> Optional[int]:
        """Get the Direct Connect port."""
        return self._direct_connect_port

    @property
    def is_connected(self) -> bool:
        """Check if the session is connected."""
        return self._connected and self._session_id is not None

    def _get_session_client(self) -> SessionClient:
        """Get or create the session client."""
        if self._session_client is None:
            self._session_client = SessionClient(self.config)
        return self._session_client

    def negotiate_session(
        self,
        host_id: Optional[str] = None,
        model_name: Optional[str] = None,
        direct_connect_required: bool = False,
        preferred_region: Optional[str] = None,
    ) -> str:
        """Negotiate a new session with the Orc.

        This handles the full session negotiation flow including:
        - Creating the session
        - Following Orc redirects (MoveToOrc)
        - Setting up Direct Connect if available
        - Connecting to the session

        Args:
            host_id: Specific host ID to connect to.
            model_name: Model name to use.
            direct_connect_required: Whether DC is required.
            preferred_region: Preferred region for host selection.

        Returns:
            The session ID.

        Raises:
            DaisiError: If negotiation fails.
            ValueError: If host_id conflicts with manager's bound host.
        """
        if self._session_id is not None:
            return self._session_id

        target_host = host_id or self.host_id
        if self.host_id and host_id and self.host_id != host_id:
            raise ValueError(
                f"SessionManager was created for host {self.host_id} "
                f"and cannot be used for host {host_id}. "
                "Create a new SessionManager instance."
            )

        logger.info("Negotiating session...")
        session_client = self._get_session_client()

        session_id, host, move_to_orc = session_client.create(
            host_id=target_host,
            model_name=model_name,
            direct_connect_required=direct_connect_required,
            preferred_region=preferred_region,
        )

        # Handle Orc redirect if needed
        if move_to_orc is not None:
            logger.info(f"Redirecting to Orc: {move_to_orc.Domain}:{move_to_orc.Port}")
            # Update config to point to new Orc
            self.config.orc_address = move_to_orc.Domain
            self.config.orc_port = move_to_orc.Port
            self.config.use_ssl = move_to_orc.RequiresSSL

            # Reset client to reconnect to new Orc
            if self._session_client:
                self._session_client.close()
                self._session_client = None

            # Retry session creation with new Orc
            session_client = self._get_session_client()
            session_id, host, move_to_orc = session_client.create(
                host_id=target_host,
                model_name=model_name,
                direct_connect_required=direct_connect_required,
                preferred_region=preferred_region,
            )

            if move_to_orc is not None:
                raise DaisiError("Multiple Orc redirects not supported")

        if host is None:
            raise DaisiError("Session created but no host assigned")

        self._session_id = session_id
        self._host = host
        logger.info(f"Session created: {session_id}")

        # Connect to the session
        logger.info("Connecting to session...")
        conn_id, has_capacity, already_connected = session_client.connect(session_id)

        if not has_capacity and not already_connected:
            raise DaisiError("Session has no capacity and is not already connected")

        self._connected = True

        # Setup Direct Connect if available
        if host.DirectConnect:
            logger.info("Direct Connect enabled")
            self._use_direct_connect = True
            if host.IpAddress and host.IpAddress.value:
                self._direct_connect_address = host.IpAddress.value
            self._direct_connect_port = host.Port
        else:
            logger.info("Using Fully Orchestrated Connection")

        return session_id

    async def negotiate_session_async(
        self,
        host_id: Optional[str] = None,
        model_name: Optional[str] = None,
        direct_connect_required: bool = False,
        preferred_region: Optional[str] = None,
    ) -> str:
        """Async version of negotiate_session.

        See negotiate_session for full documentation.
        """
        # For now, we just call the sync version
        # A proper async implementation would use AsyncSessionClient
        return self.negotiate_session(
            host_id=host_id,
            model_name=model_name,
            direct_connect_required=direct_connect_required,
            preferred_region=preferred_region,
        )

    def close(self) -> None:
        """Close the session."""
        if self._session_id and self._session_client:
            logger.info(f"Closing session: {self._session_id}")
            try:
                self._session_client.close_session(self._session_id)
            except DaisiError:
                pass  # Best effort close
            finally:
                self._session_id = None
                self._connected = False
                self._host = None
                self._use_direct_connect = False
                self._direct_connect_address = None
                self._direct_connect_port = None

        if self._session_client:
            self._session_client.close()
            self._session_client = None

        logger.info("Session closed")

    async def close_async(self) -> None:
        """Async version of close."""
        # For now, just call sync version
        self.close()

    def __enter__(self) -> "SessionManager":
        """Context manager entry."""
        return self

    def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
        """Context manager exit."""
        self.close()

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

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


class InferenceSessionManager(SessionManager):
    """Session manager specialized for inference operations.

    This is the Python equivalent of .NET's InferenceSessionManager,
    though in Python we don't need the generic type parameter since
    we use duck typing.
    """

    def __init__(
        self,
        config: Optional[DaisiConfig] = None,
        host_id: Optional[str] = None,
    ) -> None:
        """Initialize inference session manager.

        Args:
            config: Daisi configuration.
            host_id: Optional host ID to bind to.
        """
        super().__init__(config, host_id)
        self._inference_id: Optional[str] = None

    @property
    def inference_id(self) -> Optional[str]:
        """Get the current inference ID."""
        return self._inference_id

    @inference_id.setter
    def inference_id(self, value: Optional[str]) -> None:
        """Set the inference ID."""
        self._inference_id = value


class SettingsSessionManager(SessionManager):
    """Session manager for settings operations."""

    pass


class PeerSessionManager(SessionManager):
    """Session manager for peer operations."""

    pass
