from os import environ

import openai

CLAUDE_MODELS = ["claude-sonnet", "claude-opus"]
CLAUDE_MAP = {
    "claude-sonnet-4": "claude-sonnet-4-20250514",
    "claude-sonnet-3-7": "claude-3-7-sonnet-20250219",
    "claude-opus-4": "claude-opus-4-20250514",
    "claude-opus-4-1": "claude-opus-4-20250805",
}
GEMINI_MODELS = ["gemini-2.5-flash", "gemini-2.5-pro", "gemini-2.0-flash"]

GPT_MAP = {
    "gpt-4.1": "gpt-4.1-2025-04-14",
    "gpt-4o-mini": "gpt-4o-mini",
    "gpt-4o": "gpt-4o-2024-11-20",
    "gpt-5": "gpt-5",
    "gpt-5-mini": "gpt-5-mini",
    "gpt-5-high": "gpt-5",
}

def construct_client(model):
    # gpt calls do not need a client
    if "gpt" in model:
        return GPT_MAP[model] 

    elif any([model_str in model for model_str in CLAUDE_MODELS]):
        return AnthropicClient(model_name=CLAUDE_MAP[model])

    elif any([model_str == model for model_str in GEMINI_MODELS]):
        return GoogleAIClient(model_name=model, api_key=environ["GEMINI_API_KEY"], use_floodgate=True, format_=None)

    else:
        raise NotImplementedError(f"{model} needs implementation")

def openai_call(model, prompt):
    completion = openai.chat.completions.create(
        model=GPT_MAP[model],
        messages=[
            {"role": "user", "content": prompt}
        ],
    )
    output_text = completion.choices[0].message.content
    return output_text

def openai_reasoning_call(model, prompt):
    print(prompt)
    completion = openai.chat.completions.create(
        model=GPT_MAP[model],
        messages=[
            {"role": "user", "content": 'Formatting reenabled\n' + prompt}
        ],
        reasoning_effort='high'
    )
    output_text = completion.choices[0].message.content
    return output_text


def generate_answer(model: str, prompt: str, client):
    if "high" in model:
        return openai_reasoning_call(model, prompt)

    if "gpt" in model:
        return openai_call(model, prompt)

    # seperating branches in case different call parsing
    elif any([model_str in model for model_str in CLAUDE_MODELS]):
        return xxx_call(prompt, client)

    elif any([model_str in model for model_str in GEMINI_MODELS]):
        return xxx_call(prompt, client)

    else:
        raise NotImplementedError(f"{model} needs implementation")
