from collections.abc import Sequence
from itertools import batched
from typing import Any

from openai import APIConnectionError, APITimeoutError, InternalServerError, OpenAI, RateLimitError
from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_random_exponential

from techdocs_graphrag.config import Settings

EMBEDDING_BATCH_SIZE = 96
TRANSIENT_ERRORS = (APIConnectionError, APITimeoutError, InternalServerError, RateLimitError)

with_retry = retry(
    retry=retry_if_exception_type(TRANSIENT_ERRORS),
    wait=wait_random_exponential(min=1, max=30),
    stop=stop_after_attempt(5),
    reraise=True,
)


class OpenAIGateway:
    def __init__(self, settings: Settings) -> None:
        self._client = OpenAI(api_key=settings.openai_api_key.get_secret_value())
        self._chat_model = settings.chat_model
        self._chat_temperature = settings.chat_temperature
        self._embedding_model = settings.embedding_model
        self._embedding_dimensions = settings.embedding_dimensions

    @with_retry
    def complete(self, system_prompt: str, user_prompt: str, *, json_mode: bool = False) -> str:
        options: dict[str, Any] = {}
        if json_mode:
            options["response_format"] = {"type": "json_object"}
        if self._chat_temperature is not None:
            options["temperature"] = self._chat_temperature
        response = self._client.chat.completions.create(
            model=self._chat_model,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            **options,
        )
        return response.choices[0].message.content or ""

    def embed(self, texts: Sequence[str]) -> list[list[float]]:
        return [
            vector
            for batch in batched(texts, EMBEDDING_BATCH_SIZE)
            for vector in self._embed_batch(batch)
        ]

    @with_retry
    def _embed_batch(self, texts: Sequence[str]) -> list[list[float]]:
        response = self._client.embeddings.create(
            model=self._embedding_model,
            input=list(texts),
            dimensions=self._embedding_dimensions,
        )
        return [item.embedding for item in sorted(response.data, key=lambda item: item.index)]