from __future__ import annotations

from typing import Any, List, Sequence

from ...core.protocols import ChatModel, EmbeddingModel
from ...core.types import ChatResponse, EmbeddingResponse, Message
from .facade import ChatOpenAI, OpenAIEmbeddings, _extract_embeddings


class OpenAIChatModelAdapter(ChatModel):
    def __init__(self, inner: ChatOpenAI) -> None:
        self.inner = inner

    def invoke(
        self,
        messages: Sequence[Message],
        **kwargs: Any,
    ) -> ChatResponse:
        return self.inner.invoke_messages(list(messages), **kwargs)


class OpenAIEmbeddingModelAdapter(EmbeddingModel):
    def __init__(self, inner: OpenAIEmbeddings) -> None:
        self.inner = inner

    def embed_query(self, text: str, **kwargs: Any) -> List[float]:
        return self.inner.embed_query(text, **kwargs)

    def embed_documents(self, texts: Sequence[str], **kwargs: Any) -> List[List[float]]:
        return self.inner.embed_documents(list(texts), **kwargs)

    def invoke(self, input: Any, **kwargs: Any) -> EmbeddingResponse:
        raw = self.inner.invoke(input, **kwargs)
        vectors = _extract_embeddings(raw)
        return EmbeddingResponse(vectors=vectors, raw=raw)
