import os
import requests

GROQ_API = "https://api.groq.com/openai/v1/chat/completions"
GROQ_KEY = os.environ.get("GROQ_API_KEY")

MODEL = "llama-3.1-8b-instant"   # 🔥 LIVE + SUPPORTED

def generate_code(task):
    if not GROQ_KEY:
        raise RuntimeError("GROQ_API_KEY not set")

    headers = {
        "Authorization": f"Bearer {GROQ_KEY}",
        "Content-Type": "application/json"
    }

    prompt = f"""
You are a senior Python engineer.

Write safe Python code to accomplish this task:

{task}

Rules:
- No OS commands
- No file deletion
- No network calls
- Must define a function
- Must assign final output to variable named `result`
- Must not use input()
- Return only Python code
"""

    payload = {
        "model": MODEL,
        "messages": [
            {"role": "user", "content": prompt}
        ],
        "max_tokens": 300,
        "temperature": 0.2,
        "top_p": 1.0
    }

    r = requests.post(GROQ_API, headers=headers, json=payload, timeout=60)

    if r.status_code != 200:
        raise RuntimeError(f"Groq error {r.status_code}: {r.text}")

    data = r.json()
    text = data["choices"][0]["message"]["content"]

    if "```" in text:
        text = text.split("```")[-1]

    return text.strip()