from typing import Optional

from openai import AuthenticationError, OpenAI

from ask.config import config
from ask.constants import OPENAI_BASE_URL, OPENAI_DEFAULT_MODEL, PROMPT
from ask.llms.inference_provider_base import InferenceProvider
from ask.llms.types import OptionsResponse


class OpenAIProvider(InferenceProvider):
    AUTH_ERROR_MESSAGE = (
        "Error: There was an error with your OpenAI API key. You can change it by running `ask --setup`."
    )

    def __init__(self):
        if not config.openai_api_key:
            raise ValueError("OPENAI_API_KEY must be set. Try running `ask --setup`.")

        self.client = OpenAI(base_url=OPENAI_BASE_URL, api_key=config.openai_api_key)
        self.model = config.openai_model or OPENAI_DEFAULT_MODEL

    def get_options(self, prompt: str, context: str) -> Optional[OptionsResponse]:
        try:
            assembled_prompt = PROMPT.format(prompt=prompt, context=context)
            response = self.client.beta.chat.completions.parse(
                model=self.model,
                messages=[{"role": "user", "content": assembled_prompt}],
                response_format=OptionsResponse,
            )
            return response.choices[0].message.parsed
        except AuthenticationError:
            print(self.AUTH_ERROR_MESSAGE)
            return None
