from flask import Flask, Response, request, send_file
import os, urllib.parse, requests

app = Flask(__name__)

URL_TEMPLATE = os.environ.get("FLUX")
if not URL_TEMPLATE:
    raise RuntimeError("Missing FLUX env variable")

@app.route("/", methods=["GET"])
def generate_or_default():
    # If no query string present, serve default.png
    if not request.query_string:
        # Make sure default.png is in the working dir (e.g. next to app.py)
        return send_file(
            "default.png",
            mimetype="image/png",
            as_attachment=False,
            download_name="default.png"
        )

    # Otherwise, generate via Pollinations
    prompt = request.args.get(
        "prompt",
        "a glowing board with the word 'OptimFLUX'"
    )
    width  = request.args.get("width", "1024")
    height = request.args.get("height", "1024")
    seed   = request.args.get("seed", "0")

    encoded_prompt = urllib.parse.quote(prompt, safe="")
    url = (
        URL_TEMPLATE
        .replace("[prompt]", encoded_prompt)
        .replace("[w]", width)
        .replace("[h]", height)
        .replace("[seed]", seed)
    )

    resp = requests.get(url, stream=True)
    if resp.status_code != 200:
        return Response(
            f"Failed to fetch image. Upstream status: {resp.status_code}",
            status=502
        )

    return Response(resp.content, content_type=resp.headers.get("Content-Type"))

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