# 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, Hugging Face). 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:
            `ModelResponse`: A unified response object containing text and metadata.
        """
        pass

    async def run(self, prompt: str, **kwargs) -> "ModelResponse":
        r"""
        High-level ergonomic alias for `generate_content_async`.

        Args:
            prompt (`str`): The input prompt.
            **kwargs: Additional generation parameters.

        Returns:
            `ModelResponse`: The model's response.
        """
        return await self.generate_content_async(prompt, **kwargs)

    @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 (for non-async environments).
        """
        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"]
