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

app = Flask(__name__)


URL_TEMPLATE = os.environ.get("FLUX")
if URL_TEMPLATE is None:
    raise RuntimeError("ERR-SCRT")

@app.route("/generate", methods=["GET"])
def generate_image():
    """
    Query params:
      - prompt (string, required)
      - width  (int, optional, default=512)
      - height (int, optional, default=512)
      - seed   (int, optional, default=0)
    """
    prompt = request.args.get("prompt", "").strip()
    if not prompt:
        return Response("Error: 'prompt' is required", status=400)

    # Defaults
    width  = request.args.get("width", "512")
    height = request.args.get("height", "512")
    seed   = request.args.get("seed", "0")

    # URL‐encode the prompt
    encoded_prompt = urllib.parse.quote(prompt, safe="")

    # Build the actual Pollinations URL
    url = URL_TEMPLATE.replace("[prompt]", encoded_prompt) \
                      .replace("[w]", width) \
                      .replace("[h]", height) \
                      .replace("[seed]", seed)

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

    # Forward the content‐type header (likely "image/png" or "image/jpeg")
    content_type = resp.headers.get("Content-Type", "application/octet-stream")
    return Response(resp.content, content_type=content_type)

if __name__ == "__main__":
    # If you run `python app.py` locally, this will start Flask’s dev server.
    app.run(host="0.0.0.0", port=7860)