import time
import json

from openai import AsyncOpenAI

from animuz_core.genai.base import BaseLLM
from animuz_core.utils import init_logger


LOGGER = init_logger(__name__)


class OpenAILLMClient(BaseLLM):
    def __init__(self, system_prompt: str = None, model: str = "gpt-4o", max_token: int = 1000) -> None:
        self.client = AsyncOpenAI()
        self.system_prompt = system_prompt
        self.model = model
        self.max_token = max_token
        self.temperature = 0.0
        self.message_history = []

    async def get_reply(self, prompt: str):
        messages = []
        if self.system_prompt:
            messages.append({"role": "system", "content": self.system_prompt})
        messages.append({"role": "user", "content": prompt})

        start = time.time()
        res = await self.client.chat.completions.create(
            model=self.model,
            max_tokens=self.max_token,
            temperature=self.temperature,
            messages=messages,
        )
        LOGGER.debug(json.dumps({
            "timestamp": time.strftime('[%Y-%b-%d %H:%M:%S %z]'),
            "message": "OpenAI LLM latency",
            "latency": round(time.time() - start, 4),
            "exception": None,
        }))
        return res.choices[0].message.content
