import os
import requests
import gradio as gr

HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN:
    raise Exception("HF_TOKEN is missing.")

API_URL = "https://api-inference.huggingface.co/models/castorini/t5-base-canard"
HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"}

def rewrite(queries_text):
    if not queries_text.strip():
        return "Enter queries (one per line)."
    queries = [q.strip() for q in queries_text.split("\n") if q.strip()]
    results = []
    for q in queries:
        payload = {"inputs": q}
        try:
            r = requests.post(API_URL, headers=HEADERS, json=payload, timeout=30)
            output = r.json()
            results.append(output[0]["generated_text"] if isinstance(output, list) and len(output) > 0 and "generated_text" in output[0] else f"Error: {output}")
        except Exception as e:
            results.append(f"Error: {str(e)}")
    return "\n\n".join([f"Original: {q}\nRewritten: {r}\n---" for q, r in zip(queries, results)])

demo = gr.Interface(
    fn=rewrite,
    inputs=gr.Textbox(label="Queries (one per line)", lines=5, placeholder="What is AI?\nHow does it work?"),
    outputs=gr.Textbox(label="Rewritten Results", lines=10),
    title="CANARD Rewrite API",
    description="Rewrites queries using T5-base-canard. API: POST /api/rewrite/ with {'data': ['query1', ...]}"
)

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