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

app = Flask(__name__, static_folder="static")

# Your Gradio share link (no trailing slash)
GRADIO_SERVER_URL = os.getenv("GRADIO_SERVER_URL", "https://e35e92eb80591306dc.gradio.live").rstrip("/")


@app.route("/")
def index():
    return send_from_directory(".", "index.html")


@app.route("/static/<path:filename>")
def static_files(filename):
    return send_from_directory(app.static_folder, filename)


@app.route("/generate", methods=["POST"])
def generate():
    body = request.get_json(force=True) or {}
    prompt = str(body.get("prompt", "")).strip()
    if not prompt:
        return jsonify({"error": "prompt is required"}), 400

    # Optional generation params with sensible defaults
    max_new_tokens  = int(body.get("max_new_tokens", 200))
    temperature     = float(body.get("temperature", 0.9))
    top_p           = float(body.get("top_p", 0.95))
    rep_penalty     = float(body.get("repetition_penalty", 1.2))

    try:
        resp = requests.post(
            f"{GRADIO_SERVER_URL}/gradio_api/run/predict",
            json={
                "data": [prompt, max_new_tokens, temperature, top_p, rep_penalty],
                "fn_index": 0,
            },
            timeout=60,
        )
        resp.raise_for_status()
        generated_text = resp.json()["data"][0]

    except requests.exceptions.Timeout:
        return jsonify({"error": "Inference timed out"}), 504
    except requests.exceptions.RequestException as e:
        return jsonify({"error": "Inference failed", "details": str(e)}), 502
    except (KeyError, IndexError) as e:
        return jsonify({"error": "Unexpected response from model server", "details": str(e)}), 502

    return jsonify({"generated_text": generated_text, "prompt": prompt})


if __name__ == "__main__":
    port = int(os.getenv("PORT", 7890))
    app.run(host="0.0.0.0", port=port)