import io
import logging
import os
import threading
import uuid

import torch
from flask import Flask, jsonify, request, send_file, redirect
from flask_cors import CORS

os.environ["OMP_NUM_THREADS"] = "4"
os.environ["MKL_NUM_THREADS"] = "4"
torch.set_num_threads(4)
torch.set_num_interop_threads(1)

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
logger = logging.getLogger("sd-turbo-server")

DEFAULT_STEPS = int(os.environ.get("DEFAULT_STEPS", "4"))
DEFAULT_GUIDANCE = float(os.environ.get("DEFAULT_GUIDANCE", "1.0"))
DEFAULT_WIDTH = 768
DEFAULT_HEIGHT = 768

app = Flask(__name__)
CORS(app)

_pipeline = None
_pipeline_lock = threading.Lock()
_pipeline_load_error = None


def load_pipeline():
    global _pipeline, _pipeline_load_error
    from diffusers import AutoPipelineForText2Image

    model_dir = os.environ.get("MODEL_LOCAL_DIR", "/app/models/sd-turbo")
    logger.info("Loading pipeline from %s", model_dir)
    try:
        pipeline = AutoPipelineForText2Image.from_pretrained(
            model_dir, torch_dtype=torch.float32, safety_checker=None
        )
        pipeline.to("cpu")
        pipeline.unet = torch.quantization.quantize_dynamic(
            pipeline.unet, {torch.nn.Linear}, dtype=torch.qint8
        )
        if hasattr(pipeline, "text_encoder"):
            pipeline.text_encoder = torch.quantization.quantize_dynamic(
                pipeline.text_encoder, {torch.nn.Linear}, dtype=torch.qint8
            )
        pipeline.set_progress_bar_config(disable=True)
        _pipeline = pipeline
        logger.info("Pipeline loaded and quantized successfully.")
    except Exception as e:
        _pipeline_load_error = str(e)
        logger.error("Pipeline load failed: %s", e)


def run_generation(prompt, width, height, steps, guidance):
    with _pipeline_lock:
        result = _pipeline(
            prompt=prompt,
            num_inference_steps=steps,
            guidance_scale=guidance,
            width=width,
            height=height,
        )
        return result.images[0]


@app.route("/")
def index():
    return "SD Turbo Server - Use /image or /v1/images/generations"


@app.route("/health")
def health():
    return jsonify({"status": "ok", "model": "sd-turbo"})


@app.route("/v1/models")
def list_models():
    return jsonify({
        "object": "list",
        "data": [{"id": "sd-turbo", "object": "model", "owned_by": "stabilityai"}]
    })


@app.route("/v1/images/generations", methods=["POST"])
def images_generations():
    """OpenAI-compatible endpoint that returns a direct image URL"""
    if not request.is_json:
        return jsonify({"error": {"message": "JSON required"}}), 400

    data = request.get_json()
    prompt = data.get("prompt")
    if not prompt:
        return jsonify({"error": {"message": "prompt is required"}}), 400

    size = data.get("size", "768x768")
    n = data.get("n", 1)
    steps = int(data.get("num_inference_steps", DEFAULT_STEPS))
    guidance = float(data.get("guidance_scale", DEFAULT_GUIDANCE))

    if steps > 6: steps = 6
    if guidance > 2.0: guidance = 1.5

    width = DEFAULT_WIDTH
    height = DEFAULT_HEIGHT
    if "x" in size:
        try:
            w, h = map(int, size.split("x"))
            width, height = min(w, 768), min(h, 768)
        except:
            pass

    request_id = uuid.uuid4().hex[:8]
    logger.info(f"[{request_id}] OpenAI gen - prompt={prompt[:60]}... size={width}x{height}")

    try:
        image = run_generation(prompt, width, height, steps, guidance)
        
        # Save to a temporary in-memory route so we can return a URL
        img_id = uuid.uuid4().hex[:12]
        # For simplicity we return a direct /image URL with the same prompt
        image_url = f"https://cloudunity-sdturbolumi.hf.space/image?prompt={prompt.replace(' ', '+')}&width={width}&height={height}&steps={steps}&guidance={guidance}"
        
        return jsonify({
            "created": int(time.time()),
            "data": [{"url": image_url}]
        })
    except Exception as e:
        logger.exception("Generation failed")
        return jsonify({"error": {"message": str(e)}}), 500


@app.route("/image", methods=["GET"])
def generate_image():
    prompt = request.args.get("prompt", "").replace("+", " ")
    if not prompt:
        return "Missing prompt", 400

    try:
        width = int(request.args.get("width", DEFAULT_WIDTH))
        height = int(request.args.get("height", DEFAULT_HEIGHT))
        steps = int(request.args.get("steps", DEFAULT_STEPS))
        guidance = float(request.args.get("guidance", DEFAULT_GUIDANCE))
    except ValueError:
        return "Invalid parameters", 400

    if width > 768: width = 768
    if height > 768: height = 768
    if steps > 6: steps = 6
    if guidance > 2.0: guidance = 1.0

    try:
        img = run_generation(prompt, width, height, steps, guidance)
        buf = io.BytesIO()
        img.save(buf, format="PNG")
        buf.seek(0)
        return send_file(buf, mimetype="image/png")
    except Exception as e:
        logger.exception("Image generation failed")
        return str(e), 500


if __name__ == "__main__":
    load_pipeline()
    port = int(os.environ.get("PORT", "7860"))
    app.run(host="0.0.0.0", port=port, threaded=True)