"""
Session client for Daisi SDK.
"""

from typing import Optional

import grpc

from google.protobuf.wrappers_pb2 import StringValue

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

from ..base_client import BaseClient
from ..config import DaisiConfig
from ..exceptions import DaisiError


class SessionClient(BaseClient):
    """Client for session management operations."""

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

        Args:
            config: Daisi configuration.
        """
        super().__init__(config)
        self._stub: Optional[Sessions_pb2_grpc.SessionsProtoStub] = None
        self.session_id: Optional[str] = None

    @property
    def stub(self) -> Sessions_pb2_grpc.SessionsProtoStub:
        """Get or create the gRPC stub."""
        if self._stub is None:
            self._stub = Sessions_pb2_grpc.SessionsProtoStub(self._get_channel())
        return self._stub

    def create(
        self,
        model_name: Optional[str] = None,
        direct_connect_required: bool = False,
        preferred_host_names: Optional[list[str]] = None,
        preferred_region: Optional[str] = None,
        network_name: Optional[str] = None,
        host_id: Optional[str] = None,
    ) -> tuple[
        str, Optional[HostModels_pb2.Host], Optional[OrcModels_pb2.Orchestrator]
    ]:
        """Create a new session.

        Args:
            model_name: Name of the model to use.
            direct_connect_required: Whether direct connection is required.
            preferred_host_names: List of preferred host names.
            preferred_region: Preferred region for the host.
            network_name: Network name to use.
            host_id: Specific host ID to connect to.

        Returns:
            Tuple of (session_id, host, move_to_orc).
            If move_to_orc is set, client should reconnect to that orchestrator.

        Raises:
            DaisiError: If session creation fails.
        """
        request = SessionModels_pb2.CreateSessionRequest(
            ModelName=model_name or "",
            DirectConnectRequired=direct_connect_required,
        )

        if preferred_host_names:
            for name in preferred_host_names:
                request.PreferredHostNames.append(StringValue(value=name))

        if preferred_region:
            request.PreferredRegion.value = preferred_region

        # Use provided network_name or fall back to config
        effective_network = network_name or self.config.network_name
        if effective_network:
            request.NetworkName.value = effective_network

        if host_id:
            request.HostId.value = host_id

        try:
            response = self.stub.Create(request, metadata=self._get_metadata())
            self.session_id = response.Id

            host = response.Host if response.HasField("Host") else None
            move_to_orc = response.MoveToOrc if response.HasField("MoveToOrc") else None

            # Note: Success=False may still return a valid session ID
            # Only fail if no session ID and no redirect
            if not response.Id and not move_to_orc:
                raise DaisiError("Session creation failed: no session ID returned")

            return response.Id, host, move_to_orc
        except grpc.RpcError as e:
            raise DaisiError(f"Failed to create session: {e.details()}") from e

    def claim(self, session_id: str) -> tuple[bool, Optional[str]]:
        """Claim an existing session.

        Args:
            session_id: ID of the session to claim.

        Returns:
            Tuple of (success, model_name).

        Raises:
            DaisiError: If claim fails.
        """
        request = SessionModels_pb2.ClaimSessionRequest(Id=session_id)

        try:
            response = self.stub.Claim(request, metadata=self._get_metadata())
            model_name = response.ModelName.value if response.ModelName.value else None
            return response.Success, model_name
        except grpc.RpcError as e:
            raise DaisiError(f"Failed to claim session: {e.details()}") from e

    def close_session(self, session_id: Optional[str] = None) -> bool:
        """Close a session.

        Args:
            session_id: ID of the session to close. Uses current session if not provided.

        Returns:
            True if close was successful.

        Raises:
            DaisiError: If close fails.
        """
        sid = session_id or self.session_id
        if not sid:
            raise DaisiError("No session ID provided or available")

        request = SessionModels_pb2.CloseSessionRequest(Id=sid)

        try:
            response = self.stub.Close(request, metadata=self._get_metadata())
            if sid == self.session_id:
                self.session_id = None
            return response.Success
        except grpc.RpcError as e:
            raise DaisiError(f"Failed to close session: {e.details()}") from e

    def connect(self, session_id: Optional[str] = None) -> tuple[str, bool, bool]:
        """Connect to a session.

        Args:
            session_id: ID of the session to connect to. Uses current session if not provided.

        Returns:
            Tuple of (session_id, has_capacity, already_connected).

        Raises:
            DaisiError: If connect fails.
        """
        sid = session_id or self.session_id
        if not sid:
            raise DaisiError("No session ID provided or available")

        request = SessionModels_pb2.ConnectRequest(SessionId=sid)

        try:
            response = self.stub.Connect(request, metadata=self._get_metadata())
            return response.Id, response.HasCapacity, response.AlreadyConnected
        except grpc.RpcError as e:
            raise DaisiError(f"Failed to connect to session: {e.details()}") from e
