import gradio as gr
import requests
import os

# Configuration
MODEL_NAME = "0dAI/0dAI-8x7b-0761"
HF_TOKEN = os.environ.get("HF_TOKEN", "")

def query_model(prompt, max_tokens=256, temperature=0.7):
    """Query the model using Hugging Face Inference API"""
    
    # Check if token is available
    if not HF_TOKEN:
        return "❌ Error: Hugging Face token not found. Please add HF_TOKEN in Space settings → Repository secrets."
    
    try:
        API_URL = f"https://api-inference.huggingface.co/models/{MODEL_NAME}"
        headers = {"Authorization": f"Bearer {HF_TOKEN}"}
        
        payload = {
            "inputs": prompt,
            "parameters": {
                "max_new_tokens": max_tokens,
                "temperature": temperature,
                "do_sample": True,
                "return_full_text": False
            }
        }
        
        response = requests.post(API_URL, headers=headers, json=payload)
        
        if response.status_code == 200:
            result = response.json()
            if isinstance(result, list) and len(result) > 0:
                return result[0].get('generated_text', 'No generated text found')
            else:
                return f"Unexpected response format: {result}"
        elif response.status_code == 401:
            return "❌ Error: Invalid Hugging Face token. Please check your HF_TOKEN in Space settings."
        elif response.status_code == 503:
            # Model is loading
            return "🔄 Model is loading, please wait 20-30 seconds and try again..."
        else:
            error_msg = response.json().get('error', 'Unknown error')
            return f"❌ API Error ({response.status_code}): {error_msg}"
            
    except Exception as e:
        return f"❌ Request failed: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="0dAI 8x7B Model Demo", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🤖 0dAI 8x7B Model Demo
    
    This Space uses the Hugging Face Inference API to run the [0dAI/0dAI-8x7b-0761](https://huggingface.co/0dAI/0dAI-8x7b-0761) model.
    
    **Note:** Make sure you've added your Hugging Face token in Space settings!
    """)
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="Input Prompt",
                lines=4,
                placeholder="Enter your prompt here...",
                value="Explain quantum computing in simple terms:"
            )
            
            with gr.Row():
                max_tokens = gr.Slider(
                    minimum=50,
                    maximum=1024,
                    value=256,
                    step=10,
                    label="Max New Tokens"
                )
                temperature = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.7,
                    step=0.1,
                    label="Temperature"
                )
            
            generate_btn = gr.Button("🚀 Generate", variant="primary")
        
        with gr.Column():
            output = gr.Textbox(
                label="Generated Text",
                lines=8,
                show_copy_button=True
            )
    
    # Examples
    gr.Examples(
        examples=[
            ["Explain quantum computing in simple terms:"],
            ["Write a short story about a robot learning to paint:"],
            ["What are the benefits of renewable energy?"],
            ["How does machine learning work?"],
            ["Write a Python function to calculate fibonacci sequence:"]
        ],
        inputs=prompt_input,
        label="Click any example to try:"
    )
    
    # Status information
    gr.Markdown("""
    ### 🔧 Setup Instructions:
    1. Get your Hugging Face token from [Settings → Tokens](https://huggingface.co/settings/tokens)
    2. Add it to this Space in **Settings → Repository secrets** as `HF_TOKEN`
    3. The model may take 20-30 seconds to load on first request
    """)
    
    generate_btn.click(
        fn=query_model,
        inputs=[prompt_input, max_tokens, temperature],
        outputs=output
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)