"""Base LLM provider interface"""

from abc import ABC, abstractmethod
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Type

from pydantic import BaseModel

from axm.core.types import Message


class LLMProvider(ABC):
    """Base class for LLM providers"""

    @abstractmethod
    def generate(
        self,
        messages: List[Message],
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
        tools: Optional[List[Dict[str, Any]]] = None,
        response_format: Optional[Type[BaseModel]] = None,
        **kwargs: Any,
    ) -> Message:
        """Generate a response from the LLM"""
        pass

    @abstractmethod
    async def agenerate(
        self,
        messages: List[Message],
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
        tools: Optional[List[Dict[str, Any]]] = None,
        response_format: Optional[Type[BaseModel]] = None,
        **kwargs: Any,
    ) -> Message:
        """Async generate a response from the LLM"""
        pass

    @abstractmethod
    def stream(
        self,
        messages: List[Message],
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
        **kwargs: Any,
    ) -> Iterator[str]:
        """Stream response chunks from the LLM"""
        pass

    @abstractmethod
    def astream(
        self,
        messages: List[Message],
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
        **kwargs: Any,
    ) -> AsyncIterator[str]:
        """Async stream response chunks from the LLM"""
        # This should be implemented as an async generator
        # async def astream(...) -> AsyncIterator[str]:
        #     yield chunk
        pass
