import io
import logging
import os
import threading
import time
import uuid

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

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

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

# ====================== CONFIG ======================
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

MODEL_LOCAL_DIR = os.environ.get("MODEL_LOCAL_DIR", "/app/models/sd-turbo")

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

    logger.info("Loading pipeline from %s", MODEL_LOCAL_DIR)
    try:
        pipeline = AutoPipelineForText2Image.from_pretrained(
            MODEL_LOCAL_DIR,
            torch_dtype=torch.float32,
            safety_checker=None,
        )
        pipeline.to("cpu")

        # Dynamic quantization for CPU speed
        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("Failed to load pipeline: %s", e)


def run_generation(prompt, width, height, steps, guidance):
    if _pipeline is None:
        raise RuntimeError("Pipeline not initialized.")

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


def truncate_prompt(prompt, max_tokens=72):
    """Cuts very long prompts to stay under CLIP's 77 token limit"""
    words = prompt.split()
    if len(words) <= max_tokens:
        return prompt
    logger.warning(f"Prompt was truncated from {len(words)} to {max_tokens} tokens")
    return " ".join(words[:max_tokens])


# ====================== ROUTES ======================

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


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


@app.route("/v1/images/generations", methods=["POST"])
def images_generations():
    if not request.is_json:
        return jsonify({"error": {"message": "Request body must be JSON"}}), 400

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

    # Truncate long prompts to prevent CLIP errors
    prompt = truncate_prompt(prompt)

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

    # Safety caps
    if steps > 6:
        steps = 6
    if guidance > 2.0:
        guidance = 1.5

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

    request_id = uuid.uuid4().hex[:8]
    logger.info(f"[{request_id}] Generating image | steps={steps} guidance={guidance}")

    try:
        image = run_generation(prompt, width, height, steps, guidance)

        # Return a direct image URL (easy to use in HTML)
        image_url = (
            f"https://cloudunity-sdturbolumi.hf.space/image"
            f"?prompt={prompt.replace(' ', '+')}"
            f"&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")
    if not prompt:
        return "Missing prompt parameter", 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

    # Safety caps for CPU
    if width > 768: width = 768
    if height > 768: height = 768
    if steps > 6: steps = 6
    if guidance > 2.0: guidance = 1.5

    try:
        img = run_generation(prompt.replace("+", " "), 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("GET /image failed")
        return str(e), 500


# ====================== ERROR HANDLERS ======================

@app.errorhandler(404)
def not_found(e):
    return jsonify({"error": "Not found"}), 404


@app.errorhandler(500)
def internal_error(e):
    return jsonify({"error": "Internal server error"}), 500


# ====================== STARTUP ======================

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