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

app = Flask(__name__)

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

DEFAULT_PATH = os.path.join(os.getcwd(), "default.png")

@app.route("/", methods=["GET"])
def generate_or_default():
    # No query => serve default.png
    if not request.query_string:
        if os.path.exists(DEFAULT_PATH):
            return send_file(DEFAULT_PATH, mimetype="image/png")
        else:
            abort(404, description="default.png not found in container")

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

    

    # URL-encode and build
    encoded = urllib.parse.quote(prompt, safe="")
    url = (
        URL_TEMPLATE
        .replace("[prompt]", encoded)
        .replace("[w]", width)
        .replace("[h]", height)
        .replace("[seed]", seed)
    )

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

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

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