import logging
import os
import re

# TODO As we call each of these functions only once, we might want to import
# their modules only when needed, e.g., inside the function
import google.generativeai as genai

from openai import OpenAI


GEMINI_MODELS = {
    "gemini-2.0-flash" : "gemini-2.0-flash-001",
    "gemini-2.0-flash-thinking" : "gemini-2.0-flash-thinking-exp-01-21",
    "gemini-2.5-pro" : "gemini-2.5-pro-exp-03-25"
}

OPENAI_MODELS = {
    "o1" : "o1-2024-12-17",
    "o3" : "o3-2025-04-16",
    "gpt-4o" : "gpt-4o-2024-08-06",
    "gpt-4.1" : "gpt-4.1-2025-04-14"
}

DEEPSEEK_MODELS = {
    "deepseek-r1" : "deepseek-reasoner",
    "deepseek-v3" : "deepseek-chat",
}

def sanitize_llm_answer(answer):
    if "```" in answer:
        return re.search(r"```(?:\w+\n)?(.*?)```", answer, re.DOTALL).group(1)
    else:
        return answer


def get_model_id(model, available_models):
    name = available_models.get(model)
    if not name:
        logging.error(
            f"Model {name} is not supported by the chosen framework. Choose one of the following: {available_models.keys()}"
        )
        raise ValueError()
    return name


def run_gemini(model_name, prompt, temperature, top_p):
    logging.info("Retrieving GOOGLE_API_KEY")
    api_key = os.getenv("GOOGLE_API_KEY")
    genai.configure(api_key=api_key)
    generation_config = {
        "temperature": temperature,
        "top_p": top_p,
        "response_mime_type": "text/plain",
    }

    logging.info("Creating LLM model")
    model = genai.GenerativeModel(get_model_id(model_name, GEMINI_MODELS), generation_config=generation_config)
    token_count = model.count_tokens(prompt)
    logging.info(f"#tokens in prompt: {token_count}")
    logging.info("Generating answer...")
    response = model.generate_content(prompt)
    logging.info(f"Gemini's metadata:\n {response.usage_metadata}")
    logging.info("Answered generated!")
    answer = response.text
    return answer


def run_deepseek(model_name, prompt, temperature, top_p):
    api_key = os.getenv("DEEPSEEK_API_KEY")
    client = OpenAI(
        base_url="https://api.deepseek.com",
        api_key=api_key,
    )

    completion = client.chat.completions.create(
        model=get_model_id(model_name, DEEPSEEK_MODELS),
        messages=[{"role": "user", "content": prompt}],
        temperature=temperature,
        top_p=top_p,
        stream=False
    )

    logging.info("LLM Usage:")
    print(completion.usage)

    answer = completion.choices[0].message.content

    return answer


def run_openai(model_name, prompt, temperature, top_p):
    api_key = os.getenv("OPENAI_API_KEY")
    client = OpenAI(
        api_key=api_key
    )

    completion = client.chat.completions.create(
        model=get_model_id(model_name, OPENAI_MODELS),
        messages=[{"role": "user", "content": prompt}],
        temperature=temperature,
        #top_p=top_p, # OpenAI does not support top-p with o1
        stream=False
    )

    logging.info("LLM Usage:")
    print(completion.usage)

    answer = completion.choices[0].message.content

    return answer
