from abc import ABC, abstractmethod
from collections.abc import AsyncIterator

from api.schemas.chat import Message
from core.trace_context import LLMUsageRecord, get_current_trace
from core.tracing import log_llm_invocation


class BaseProvider(ABC):
    """
    All LLM providers implement this interface.
    Agents talk to this — never to a concrete provider directly.
    """

    def __init__(self, model: str) -> None:
        self.model = model

    @property
    @abstractmethod
    def name(self) -> str:
        """Provider identifier, e.g. 'ollama', 'openai'."""
        ...

    @abstractmethod
    async def complete(
        self,
        messages: list[Message],
        system_prompt: str | None = None,
    ) -> Message:
        """
        Send a list of messages and return the assistant reply.
        Implementations must return a Message with role='assistant'.
        """
        ...

    async def stream(
        self,
        messages: list[Message],
        system_prompt: str | None = None,
    ) -> AsyncIterator[str]:
        """Async generator yielding incremental token strings."""
        raise NotImplementedError(
            f"{self.__class__.__name__} does not yet support streaming."
        )
        yield  # pragma: no cover — makes this a generator for type checkers

    def _log_llm_usage(self, *, purpose: str, streaming: bool) -> None:
        trace = get_current_trace()
        if trace is None:
            return
        log_llm_invocation(
            trace,
            LLMUsageRecord(
                purpose=purpose,
                provider=self.name,
                model=self.model,
                streaming=streaming,
            ),
            component=self.__class__.__name__,
        )