import os
import requests
import json
import llm_gemini

AVAILABLE_MODELS = [
    "claude-3-haiku",
    "claude-3-sonnet",
    "claude-3-opus",
    "gemini-pro",
    "gpt-3.5-turbo",
    "gpt-4-turbo",
    "gpt-4",
]

MODEL = "claude-3-haiku"

URL = "https://api.theb.ai/v1/chat/completions"
# URL = "https://api.baizhi.ai/v1/chat/completions"
API_KEY = os.environ.get("THEB_API_KEY", "[THEB_API_KEY]")


def generate_content(prompt, model=MODEL, api_key=API_KEY):
    assert model in AVAILABLE_MODELS, f"Model {model} not available"
    if model == "gemini-pro":
        return llm_gemini.generate_content(prompt)
    payload = json.dumps({
        "model": model,
        "messages": [
            {
            "role": "user",
            "content": prompt
            }
        ],
        "stream": False
        })
    headers = {
        'Authorization': f'Bearer {api_key}',
        'Content-Type': 'application/json'
    }

    response = requests.request("POST", URL, headers=headers, data=payload, timeout=10)
    if response.status_code != 200:
        retry_count = 0
        while retry_count < 5:
            response = requests.request("POST", URL, headers=headers, data=payload, timeout=10)
            if response.status_code == 200:
                break
            else:
                retry_count += 1
        if retry_count == 5:
            raise Exception(f"Failed to get completion from LLM. Status code: {response.status_code}. Response: {response.text}")

    # print(response)

    return response.json()['choices'][0]['message']['content']


if __name__ == "__main__":
    print(generate_content("What are you?", "claude-3-haiku"))