# coding=utf-8
# Copyright 2026 XRTM Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from abc import ABC, abstractmethod
from typing import Any, AsyncIterable, Dict, List, Optional


class InferenceProvider(ABC):
    r"""
    Abstract Base Class for Large Language Model (LLM) providers.

    `InferenceProvider` standardizes the interface for interacting with various
    LLM SDKs (e.g., Gemini, OpenAI, Anthropic). This allows the rest of the
    platform to remain agnostic to the specific backend being used.
    """

    @abstractmethod
    async def generate_content_async(self, prompt: str, output_logprobs: bool = False, **kwargs) -> Any:
        r"""
        Asynchronously generates content from the LLM.

        Args:
            prompt (`str`):
                The input prompt to send to the model.
            output_logprobs (`bool`, *optional*, defaults to `False`):
                Whether to return log probabilities for the generated tokens.
            **kwargs:
                Additional provider-specific parameters (e.g., temperature, max_tokens).

        Returns:
            `Any`: A `ModelResponse` instance containing the generated text and metadata.
        """
        pass

    @property
    def supports_tools(self) -> bool:
        r"""
        Feature flag indicating if the provider supports native function calling.

        Returns:
            `bool`: `True` if the provider supports tool use, else `False`.
        """
        return True

    @abstractmethod
    def generate_content(self, prompt: str, output_logprobs: bool = False, **kwargs) -> Any:
        r"""
        Synchronously generates content from the LLM.
        """
        pass

    @abstractmethod
    def stream(self, messages: List[Dict[str, str]], **kwargs) -> AsyncIterable[Any]:
        r"""
        Opens a streaming connection to the LLM for token-by-token generation.
        """
        pass


class ModelResponse:
    r"""
    A unified wrapper for LLM responses, normalizing field access across different providers.

    Attributes:
        text (`str`):
            The primary text content generated by the model.
        raw (`Any`):
            The original, un-parsed response object from the provider's SDK.
        usage (`Dict[str, int]`):
            Token statistics (prompt, completion, and total).
        logprobs (`Optional[List[Dict[str, Any]]]`):
            Optional token-level log probabilities if requested.
    """

    def __init__(
        self,
        text: str,
        raw: Any = None,
        usage: Optional[Dict[str, int]] = None,
        logprobs: Optional[List[Dict[str, Any]]] = None,
    ):
        self.text = text
        self.raw = raw
        self.usage = usage or {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
        self.logprobs = logprobs


__all__ = ["InferenceProvider", "ModelResponse"]
