import json
import traceback
from flask import Flask, render_template, request, jsonify, send_from_directory, Response
from sd_pipeline import pipeline
from models import GenerationOptions

app = Flask(__name__)


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


@app.route("/generate", methods=["POST"])
def generate():
    data = request.get_json()

    prompt = data.get("prompt", "")
    if not prompt:
        return jsonify({"success": False, "error": "Prompt is required"}), 400

    # Parse and validate seed
    seed = data.get("seed")
    if seed is not None and seed != "":
        seed = int(seed)
    else:
        seed = None

    # Parse and clamp numeric values
    width = data.get("width")
    height = data.get("height")
    if width:
        width = max(256, min(2048, int(width) // 8 * 8))  # must be multiple of 8
    if height:
        height = max(256, min(2048, int(height) // 8 * 8))

    options = GenerationOptions(
        prompt=prompt,
        negative_prompt=data.get("negative_prompt", ""),
        seed=seed,
        steps=max(1, min(100, int(data.get("steps", 20)))),
        guidance_scale=max(1.0, min(20.0, float(data.get("guidance_scale", 7.5)))),
        count=max(1, min(10, int(data.get("count", 1)))),
        add_quality_keywords=data.get("add_quality_keywords", True),
        increment_seed=data.get("increment_seed", True),
        vary_guidance=data.get("vary_guidance", False),
        guidance_low=max(1.0, min(20.0, float(data.get("guidance_low", 5.0)))),
        guidance_high=max(1.0, min(20.0, float(data.get("guidance_high", 12.0)))),
        vary_steps=data.get("vary_steps", False),
        steps_low=max(1, min(100, int(data.get("steps_low", 20)))),
        steps_high=max(1, min(100, int(data.get("steps_high", 80)))),
        width=width,
        height=height,
    )

    def generate_events():
        try:
            for result in pipeline.generate_stream(options):
                yield f"data: {json.dumps(result.to_dict())}\n\n"
            yield f"data: {json.dumps({'done': True})}\n\n"
        except Exception as e:
            traceback.print_exc()
            yield f"data: {json.dumps({'error': str(e)})}\n\n"

    return Response(generate_events(), mimetype='text/event-stream')


@app.route("/out/<path:filename>")
def serve_image(filename):
    return send_from_directory("out", filename)


if __name__ == "__main__":
    print("Loading model on startup...")
    pipeline.load()
    print("Starting web server...")
    app.run(host="127.0.0.1", port=5000, debug=False, threaded=True)