import os
import json
from groq import Groq
import google.generativeai as genai

# ==========================
# Setup API Clients
# ==========================

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

genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
gemini_model = genai.GenerativeModel("gemini-2.5-flash")


# ==========================
# Llama (Groq)
# ==========================

def call_llama(messages, temperature=0.7):
    """
    messages = [
        {"role": "system", "content": "..."},
        {"role": "user", "content": "..."}
    ]
    """

    response = groq_client.chat.completions.create(
        model="llama-3.1-8b-instant",
        messages=messages,
        temperature=temperature
    )

    return response.choices[0].message.content


# ==========================
# Gemini
# ==========================

def call_gemini(messages, temperature=0.7):
    """
    Gemini expects combined prompt.
    We convert messages into a single string.
    """

    formatted_prompt = ""
    for m in messages:
        formatted_prompt += f"{m['role'].upper()}: {m['content']}\n"

    response = gemini_model.generate_content(
        formatted_prompt,
        generation_config={"temperature": temperature}
    )

    return response.text


# ==========================
# AI Classifier
# ==========================

def classify_prompt(prompt):

    system_prompt = """
You are an AI intent classifier.

Return ONLY valid JSON in this format:

{
  "intent": "chat | reasoning | live_data | greeting",
  "needs_search": true or false,
  "complexity": "low | medium | high"
}

Rules:
- Logical, math, proof, step explanation -> reasoning
- Current events, weather, news, price, today -> live_data
- Hello/hi/how are you -> greeting
- Everything else -> chat
"""

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

    raw = call_llama(messages, temperature=0)

    try:
        return json.loads(raw)
    except:
        return {
            "intent": "chat",
            "needs_search": False,
            "complexity": "low"
        }


# ==========================
# LLM Judge (for dual answers)
# ==========================

def judge_answers(answer1, answer2):

    judge_prompt = f"""
Choose which answer is better.

Reply ONLY with:
1
or
2

Answer 1:
{answer1}

Answer 2:
{answer2}
"""

    result = call_llama(
        [{"role": "user", "content": judge_prompt}],
        temperature=0
    ).strip()

    if result == "2":
        return 2
    return 1