import os
import logging

import httpx
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

logging.basicConfig(level=logging.INFO)


class LLMClient:
    def __init__(self):
        llama_url = os.getenv("LLAMA_SERVER_URL")
        if llama_url:
            # Tight read timeout so a wedged llama-server fails fast instead
            # of hanging for the OpenAI client's 600s default.
            read_timeout = float(os.getenv("LLAMA_TIMEOUT_SECONDS", "120"))
            self.client = OpenAI(
                base_url=llama_url,
                api_key="not-needed",
                timeout=httpx.Timeout(read_timeout, connect=10.0),
                max_retries=1,
            )
            self.model = os.getenv("LLAMA_MODEL_NAME", "llama-3.1-8b-instruct")
            self.PROVIDER = "llama_server"
            logging.info("Using llama_server as LLM backend")
        else:
            self.client = OpenAI()
            self.model = "gpt-4o-mini"
            self.PROVIDER = "openai"
            logging.info("Using OpenAI as LLM backend")

    def chat(
        self,
        prompt: str,
        system_prompt: str,
    ):
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {
                    "role": "system",
                    "content": system_prompt,
                },
                {"role": "user", "content": prompt},
            ],
        )
        return response.choices[0].message.content


if __name__ == "__main__":
    client = LLMClient()
    print(client.chat(prompt="Hello!", system_prompt="You are a helpful assistant."))