import os
from dotenv import load_dotenv
import openai
import gradio as gr

# Load secret key from environment
XAI_API_KEY = os.getenv("XAI_API_KEY")
if not XAI_API_KEY:
    raise ValueError("❌ XAI_API_KEY not found in environment.")

# OpenAI client for Grok
client = openai.OpenAI(api_key=XAI_API_KEY, base_url="https://api.x.ai/v1")

# Function to query Grok
def query_grok(prompt):
    try:
        completion = client.chat.completions.create(
            model="grok-beta",
            messages=[
                {"role": "system", "content": "You are a helpful AI assistant."},
                {"role": "user", "content": prompt},
            ],
            temperature=0.2,
            max_tokens=500
        )
        return completion.choices[0].message.content
    except Exception as e:
        return f"❌ API Error: {str(e)}"

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🚀 GROK API Test in Hugging Face Space")
    prompt_box = gr.Textbox(label="Enter your query", lines=2)
    output_box = gr.Textbox(label="Grok Response", lines=10)
    btn = gr.Button("Send")
    btn.click(query_grok, inputs=prompt_box, outputs=output_box)

demo.launch()