import os
from dotenv import load_dotenv
from groq import Groq

load_dotenv()

client = Groq(api_key=os.getenv("GROQ_API_KEY"))

MODEL = "llama-3.3-70b-versatile"


def ask_llm(retrieved_chunks, question, history=None):
    if history is None:
        history = []

    if isinstance(retrieved_chunks, list):
        context = "\n\n".join(
            [r["chunk"] for r in retrieved_chunks if isinstance(r, dict) and "chunk" in r]
        )
    else:
        context = str(retrieved_chunks)

    system_prompt = f"""You are an Enterprise AI Knowledge Assistant.

Answer ONLY using the provided context.

If the answer is not found in the context, reply exactly:
"I couldn't find that information in the uploaded document."

Context:
{context}
"""

    messages = [
        {"role": "system", "content": system_prompt}
    ]

    # Append conversation history (sanitizes messages to only include role and content for Groq API)
    for msg in history:
        if isinstance(msg, dict) and "role" in msg and "content" in msg:
            messages.append({"role": str(msg["role"]), "content": str(msg["content"])})
        elif isinstance(msg, (list, tuple)) and len(msg) == 2:
            messages.append({"role": "user", "content": str(msg[0])})
            messages.append({"role": "assistant", "content": str(msg[1])})
        elif hasattr(msg, "role") and hasattr(msg, "content"):
            messages.append({"role": str(getattr(msg, "role")), "content": str(getattr(msg, "content"))})

    # Append current user question
    messages.append({"role": "user", "content": question})

    response = client.chat.completions.create(
        model=MODEL,
        messages=messages,
        temperature=0.2,
        max_tokens=1024,
    )

    return response.choices[0].message.content