import os
import requests
from flask import Flask, request, jsonify
from flask_cors import CORS  # For cross-origin requests

app = Flask(__name__)
CORS(app)  # Enable CORS

# 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://router.huggingface.co/hf-inference/models/castorini/t5-base-canard"
HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"}

@app.route("/")
def home():
    return "CANARD Rewrite API running on HuggingFace Spaces! Use POST to / with {'queries': [...]}"

@app.route("/rewrite", methods=["POST"])  # Moved to root for HF routing compatibility
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}
        try:
            r = requests.post(API_URL, headers=HEADERS, json=payload, timeout=30)  # Add timeout
            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}")
        except requests.exceptions.RequestException as e:
            results.append(f"Request failed: {str(e)}")

    return jsonify({"results": results})

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