import os
from typing import Literal

from pydantic import SecretStr
from _utils.LLMs.LLM_class import LLM
from setup.environment import default_model
from langchain_core.messages import HumanMessage
from langchain_openai import ChatOpenAI

llm = LLM()

async def aclaude_answer(claude_client, claude_context_model, prompt):
    print("\n\nComeçou uma requisição pelo Claude")
    response = await claude_client.messages.create(
        model=claude_context_model,
        max_tokens=100,  # Máximo é 4096
        messages=[{"role": "user", "content": prompt}],
    )
    return response.content[
        0
    ].text  # O response.content é uma lista pois é passada uma lista de mensagens, e também retornado uma lista de mensagens, sendo a primeira a mais recente, que é a resposta do model


async def agpt_answer(prompt):
    api_key = os.environ.get("OPENAI_API_KEY")
    if api_key:
        gpt = ChatOpenAI(
            temperature=0,
            model="gpt-4o-mini",
            api_key=SecretStr(api_key),
            max_retries=5,
        )
        response = await gpt.ainvoke([HumanMessage(content=prompt)])
        return response.content


def gpt_answer(
    prompt,
    temperature=0,
    model=default_model,
    max_retries=5,
    shouldReturnFullResponse=False,
):
    api_key = os.environ.get("OPENAI_API_KEY")
    if api_key:
        gpt = ChatOpenAI(
            temperature=temperature,
            model=model,
            api_key=SecretStr(api_key),
            max_retries=max_retries,
        )
        response = gpt.invoke([HumanMessage(content=prompt)])
        if shouldReturnFullResponse:
            return response
        else:
            return response.content

async def agemini_answer(prompt, model: Literal["gemini-2.0-flash", "gemini-2.0-flash-lite"]="gemini-2.0-flash"):
    gemini = llm.google_gemini(model)
    resposta = await gemini.ainvoke([HumanMessage(content=prompt)])
    return resposta.content