import gradio as gr
import os
import requests

# --------------------------
# Configuration
# --------------------------
HF_ENDPOINT_URL = "https://dqptyla9qxd15452.us-east-1.aws.endpoints.huggingface.cloud"
HF_TOKEN = os.environ.get("HF_TOKEN")

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

# --------------------------
# Function to Call Endpoint
# --------------------------
def generate_text(prompt, max_length=150, temperature=0.7):
    payload = {
        "inputs": prompt,
        "parameters": {
            "max_new_tokens": max_length,
            "temperature": temperature,
            "top_k": 50,
            "top_p": 0.95,
            "repetition_penalty": 1.5
        }
    }

    try:
        response = requests.post(HF_ENDPOINT_URL, headers=headers, json=payload)
        if response.status_code == 200:
            return response.json()[0]["generated_text"]
        elif response.status_code == 404:
            return "❌ Endpoint not found. Please check your HF_ENDPOINT_URL."
        elif response.status_code == 403:
            return "❌ Forbidden: Check your HF_TOKEN permissions."
        elif response.status_code == 503:
            return "❌ Service Unavailable: The endpoint may be starting up. Please wait and try again."
        else:
            return f"❌ Error {response.status_code}: {response.text}"
    except Exception as e:
        return f"❌ Failed to connect to endpoint: {str(e)}"

# --------------------------
# Gradio UI
# --------------------------
with gr.Blocks(title="Healthelic Burmese LLM (Hosted Inference Endpoint)") as demo:
    gr.Markdown("## 🧠 Healthelic Burmese LLM")
    gr.Markdown("This app connects to a Hugging Face Inference Endpoint to generate Burmese language text.\n\nEnter a prompt in Burmese, adjust settings, and click 'Generate Text'.")

    with gr.Row():
        with gr.Column(scale=3):
            prompt = gr.Textbox(
                lines=5,
                placeholder="Enter your Burmese text here...",
                label="Input Prompt"
            )
        with gr.Column(scale=1):
            max_length = gr.Slider(
                minimum=50,
                maximum=500,
                value=150,
                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 Text")
    output = gr.Textbox(lines=10, label="Generated Output")

    generate_btn.click(
        fn=generate_text,
        inputs=[prompt, max_length, temperature],
        outputs=output
    )

    with gr.Accordion("📌 Example Prompts", open=False):
        gr.Markdown("Click on a prompt below to try it out:")
        examples = [
            ["မင်္ဂလာပါ။ ကျွန်တော်က ကိုအောင်ပါ။ ရန်ကုန်မှာနေတယ်။ ဆရာလုပ်ပါတယ်။", 150, 0.7],
            ["မြန်မာနိုင်ငံမှာ မိုးရာသီမှာ မိုးဘယ်လောက်ကျလဲ။", 150, 0.6],
            ["အောင်မြင်တဲ့ကျောင်းသူတစ်ယောက်ဖြစ်ဖို့ ဘယ်လိုလေ့ကျင့်ရမလဲ။", 200, 0.8],
            ["မြန်မာ့ယဉ်ကျေးမှုအကြောင်း ပြောပြပါ။", 150, 0.7],
            ["သင်တောင်ကြီးသွားဖူးလား။ ဘယ်မှာသွားဖူးလဲ။", 150, 0.6]
        ]
        for idx, ex in enumerate(examples):
            example_btn = gr.Button(f"Example {idx+1}")
            example_btn.click(
                lambda e=ex: (e[0], e[1], e[2]),
                inputs=[],
                outputs=[prompt, max_length, temperature]
            ).then(
                fn=generate_text,
                inputs=[prompt, max_length, temperature],
                outputs=output
            )

    gr.Markdown("""
    ---
    ### ℹ️ Troubleshooting
    - If the app says "❌ Endpoint not found", double-check your `HF_ENDPOINT_URL`.
    - If the model takes a few seconds to respond, it may be waking up from idle.
    - Ensure your Hugging Face token is added as `HF_TOKEN` under **Settings → Secrets**.
    """)

# --------------------------
# Launch the Space
# --------------------------
demo.launch(
    show_error=True,
    server_name="0.0.0.0",
    share=False
)