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

app = Flask(__name__)

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

@app.route("/", methods=["GET"])
def generate_image():
    # Default parameters
    prompt = request.args.get("prompt", "a plain simple glowing board with the word 'OptimFLUX'")
    width = request.args.get("width", "1024")
    height = request.args.get("height", "1024")
    seed = request.args.get("seed", "0")

    # Sanitize prompt
    encoded_prompt = urllib.parse.quote(prompt)

    # Fill in URL
    url = (
        URL_TEMPLATE
        .replace("[prompt]", encoded_prompt)
        .replace("[w]", width)
        .replace("[h]", height)
        .replace("[seed]", seed)
    )

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

    # Serve image
    return Response(response.content, content_type=response.headers.get("Content-Type"))