#!/usr/bin/env python

"""████
 ██    ██    Datature
   ██  ██    Powering Breakthrough AI
     ██

@File    :   base.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Base inference backend interface for local deployment.
"""

from abc import ABC, abstractmethod
from collections.abc import AsyncIterator
from dataclasses import dataclass
from typing import Any


@dataclass
class GenerationResult:
    """Result from a non-streaming generation request."""

    text: str
    prompt_tokens: int
    completion_tokens: int
    finish_reason: str = "stop"


class InferenceBackend(ABC):
    """Abstract base class for inference backends.

    Defines the interface that all inference backends must implement
    to support both streaming and non-streaming generation.
    """

    @abstractmethod
    async def generate(
        self,
        messages: list[dict] | None,
        response_format: Any | None = None,
        generation_config: dict | None = None,
        **kwargs,
    ) -> GenerationResult:
        """Generate a complete response (non-streaming).

        Args:
            messages: List of chat messages in OpenAI format
            response_format: Optional structured output format spec
            generation_config: Dictionary of generation parameters (temperature, max_tokens, etc.)
            **kwargs: Additional generation parameters

        Returns:
            GenerationResult with complete response and token counts

        """

    @abstractmethod
    async def generate_stream(
        self,
        messages: list[dict] | None,
        response_format: Any | None = None,
        generation_config: dict | None = None,
        **kwargs,
    ) -> AsyncIterator[str]:
        """Generate a streaming response.

        Args:
            messages: List of chat messages in OpenAI format
            response_format: Optional structured output format spec
            generation_config: Dictionary of generation parameters (temperature, max_tokens, etc.)
            **kwargs: Additional generation parameters

        Yields:
            Tokens as strings

        """

    @abstractmethod
    async def generate_vi_format(
        self,
        source: str | None = None,
        user_prompt: str | None = None,
        response_format: Any | None = None,
        generation_config: dict | None = None,
        **kwargs,
    ) -> GenerationResult:
        """Generate a complete response using Vi SDK format (non-streaming).

        Args:
            source: Image source path (can be None for text-only requests)
            user_prompt: Text prompt (can be None)
            response_format: Optional structured output format spec
            generation_config: Dictionary of generation parameters (temperature, max_tokens, etc.)
            **kwargs: Additional generation parameters

        Returns:
            GenerationResult with complete response and token counts

        """

    @abstractmethod
    async def generate_stream_vi_format(
        self,
        source: str | None = None,
        user_prompt: str | None = None,
        response_format: Any | None = None,
        generation_config: dict | None = None,
        **kwargs,
    ) -> AsyncIterator[str]:
        """Generate a streaming response using Vi SDK format.

        Args:
            source: Image source path (can be None for text-only requests)
            user_prompt: Text prompt (can be None)
            response_format: Optional structured output format spec
            generation_config: Dictionary of generation parameters (temperature, max_tokens, etc.)
            **kwargs: Additional generation parameters

        Yields:
            Tokens as strings

        """

    @abstractmethod
    def available_models(self) -> list[str]:
        """Get list of available model identifiers.

        Returns:
            List of model IDs that can be used in generation requests

        """
