import os
import requests
from flask import Flask, request, jsonify

app = Flask(__name__)

# Load HF API token from environment variable
HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN:
    raise Exception("HF_TOKEN is missing. Add it in Space Secrets.")

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

@app.route("/")
def home():
    return "CANARD Rewrite API running on HuggingFace Spaces!"

@app.route("/rewrite", methods=["POST"])
def rewrite():
    data = request.get_json()
    queries = data.get("queries")

    if not isinstance(queries, list):
        return jsonify({"error": "queries must be a list"}), 400

    results = []

    for q in queries:
        payload = {"inputs": q}
        r = requests.post(API_URL, headers=HEADERS, json=payload)
        output = r.json()

        if isinstance(output, list) and len(output) > 0 and "generated_text" in output[0]:
            results.append(output[0]["generated_text"])
        else:
            results.append(f"Error: {output}")

    return jsonify({"results": results})

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)