"""
Inference client for Daisi SDK.
"""

from typing import Iterator, Optional

import grpc

from protos.v1 import Inferences_pb2, Inferences_pb2_grpc
from protos.v1.models import InferenceModels_pb2

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


class InferenceClient(BaseClient):
    """Client for inference operations."""

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

        Args:
            config: Daisi configuration.
            session_id: Active session ID.
        """
        super().__init__(config)
        self._stub: Optional[Inferences_pb2_grpc.InferencesProtoStub] = None
        self.session_id = session_id
        self.inference_id: Optional[str] = None

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

    def create(
        self,
        session_id: Optional[str] = None,
        model_name: Optional[str] = None,
        initialization_prompt: Optional[str] = None,
        think_level: InferenceModels_pb2.ThinkLevels = InferenceModels_pb2.ThinkBasic,
        tool_groups: Optional[list[InferenceModels_pb2.InferenceToolGroups]] = None,
    ) -> tuple[str, str]:
        """Create a new inference session.

        Args:
            session_id: Session ID. Uses stored session_id if not provided.
            model_name: Model to use for inference.
            initialization_prompt: System/initialization prompt.
            think_level: Level of reasoning (Basic, ChainOfThought, TreeOfThought).
            tool_groups: List of tool groups to enable.

        Returns:
            Tuple of (session_id, inference_id).

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

        request = InferenceModels_pb2.CreateInferenceRequest(
            SessionId=sid,
            ModelName=model_name or "",
            InitializationPrompt=initialization_prompt or "",
            ThinkLevel=think_level,
        )

        if tool_groups:
            request.ToolGroups.extend(tool_groups)

        try:
            response = self.stub.Create(request, metadata=self._get_metadata())
            self.inference_id = response.InferenceId
            return response.SessionId, response.InferenceId
        except grpc.RpcError as e:
            raise DaisiError(f"Failed to create inference: {e.details()}") from e

    def send(
        self,
        text: str,
        session_id: Optional[str] = None,
        inference_id: Optional[str] = None,
        tool_group: Optional[InferenceModels_pb2.InferenceToolGroups] = None,
        tool_name: Optional[str] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        top_k: Optional[int] = None,
        max_tokens: Optional[int] = None,
        **kwargs: Optional[int | float | bool | list[str]],
    ) -> Iterator[InferenceModels_pb2.SendInferenceResponse]:
        """Send inference request and stream responses.

        Args:
            text: The input text to process.
            session_id: Session ID. Uses stored session_id if not provided.
            inference_id: Inference ID. Uses stored inference_id if not provided.
            tool_group: Tool group to use for this request.
            tool_name: Specific tool name to invoke.
            temperature: Sampling temperature (0.0-2.0).
            top_p: Nucleus sampling threshold.
            top_k: Top-k sampling.
            max_tokens: Maximum tokens to generate.
            **kwargs: Additional sampling parameters (repeat_penalty, seed, etc).

        Yields:
            Stream of inference responses.

        Raises:
            DaisiError: If send fails.
        """
        sid = session_id or self.session_id
        iid = inference_id or self.inference_id

        if not sid:
            raise DaisiError("No session ID provided or available")
        if not iid:
            raise DaisiError("No inference ID provided or available")

        request = InferenceModels_pb2.SendInferenceRequest(
            SessionId=sid, InferenceId=iid, Text=text
        )

        if tool_group is not None:
            request.ToolGroup = tool_group
        if tool_name:
            request.ToolName = tool_name
        if temperature is not None:
            request.Temperature = temperature
        if top_p is not None:
            request.TopP = top_p
        if top_k is not None:
            request.TopK = top_k
        if max_tokens is not None:
            request.MaxTokens = max_tokens

        for key, value in kwargs.items():
            if value is not None:
                setattr(request, key, value)

        try:
            yield from self.stub.Send(request, metadata=self._get_metadata())
        except grpc.RpcError as e:
            raise DaisiError(f"Failed to send inference: {e.details()}") from e

    def stats(
        self,
        session_id: Optional[str] = None,
        inference_id: Optional[str] = None,
    ) -> InferenceModels_pb2.InferenceStatsResponse:
        """Get statistics for an inference session.

        Args:
            session_id: Session ID. Uses stored session_id if not provided.
            inference_id: Inference ID. Uses stored inference_id if not provided.

        Returns:
            Inference statistics including token counts and compute time.

        Raises:
            DaisiError: If stats request fails.
        """
        sid = session_id or self.session_id
        iid = inference_id or self.inference_id

        if not sid:
            raise DaisiError("No session ID provided or available")
        if not iid:
            raise DaisiError("No inference ID provided or available")

        request = InferenceModels_pb2.InferenceStatsRequest(
            SessionId=sid, InferenceId=iid
        )

        try:
            return self.stub.Stats(request, metadata=self._get_metadata())
        except grpc.RpcError as e:
            raise DaisiError(f"Failed to get inference stats: {e.details()}") from e

    def close_inference(
        self,
        session_id: Optional[str] = None,
        inference_id: Optional[str] = None,
        reason: InferenceModels_pb2.InferenceCloseReasons = InferenceModels_pb2.CloseRequestedByClient,
    ) -> bool:
        """Close an inference session.

        Args:
            session_id: Session ID. Uses stored session_id if not provided.
            inference_id: Inference ID. Uses stored inference_id if not provided.
            reason: Reason for closing the inference.

        Returns:
            True if close was successful.

        Raises:
            DaisiError: If close fails.
        """
        sid = session_id or self.session_id
        iid = inference_id or self.inference_id

        if not sid:
            raise DaisiError("No session ID provided or available")
        if not iid:
            raise DaisiError("No inference ID provided or available")

        request = InferenceModels_pb2.CloseInferenceRequest(
            SessionId=sid, InferenceId=iid, Reason=reason
        )

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