import os
import sys
import subprocess
import base64
import io
import traceback

# Install latest diffusers (supports GLM-Image) + dependencies
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "--upgrade",
    "diffusers", "flask", "accelerate", "sentencepiece", "protobuf", "transformers", "huggingface_hub"])

import torch
from flask import Flask, request, jsonify

print(f"[image] torch version: {torch.__version__}", flush=True)
print(f"[image] CUDA available: {torch.cuda.is_available()}", flush=True)
if torch.cuda.is_available():
    print(f"[image] GPU: {torch.cuda.get_device_name(0)}", flush=True)
    print(f"[image] VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB", flush=True)

app = Flask(__name__)
pipe = None

def load_model():
    global pipe
    from diffusers.pipelines.glm_image import GlmImagePipeline
    print("[image] Loading GLM-Image pipeline...", flush=True)
    pipe = GlmImagePipeline.from_pretrained(
        "zai-org/GLM-Image",
        torch_dtype=torch.bfloat16,
    )
    pipe.to("cuda")
    print("[image] Model loaded on CUDA!", flush=True)

@app.route("/health", methods=["GET"])
def health():
    if pipe is not None:
        return jsonify({"status": "healthy"}), 200
    return jsonify({"status": "loading"}), 503

@app.route("/", methods=["POST"])
def generate():
    try:
        data = request.get_json(force=True)
        # Support both HF format (inputs) and backend format (prompt)
        prompt = data.get("inputs", "") or data.get("prompt", "")
        params = data.get("parameters", {})

        width = int(params.get("width", 1024))
        height = int(params.get("height", 1024))
        steps = int(params.get("num_inference_steps", 30))
        guidance = float(params.get("guidance_scale", 1.5))

        # GLM-Image requires dimensions divisible by 32
        width = (width // 32) * 32
        height = (height // 32) * 32

        print(f"[image] Generating: {prompt[:100]}", flush=True)
        print(f"[image] Params: {width}x{height}, steps={steps}, guidance={guidance}", flush=True)

        image = pipe(
            prompt=prompt,
            height=height,
            width=width,
            num_inference_steps=steps,
            guidance_scale=guidance,
        ).images[0]

        buf = io.BytesIO()
        image.save(buf, format="PNG")
        img_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
        print(f"[image] Done, size: {len(buf.getvalue())} bytes", flush=True)

        # Return in both formats - HF format and backend format
        return jsonify({
            "image": img_b64,
            "format": "png",
            "success": True,
            "image_base64": img_b64,
            "model": "glm-image",
        })
    except Exception as e:
        tb = traceback.format_exc()
        print(f"[image] ERROR: {e}\n{tb}", flush=True)
        return jsonify({"error": str(e), "traceback": tb, "success": False}), 500

# Also add /api/generate route for the backend's expected format
@app.route("/api/generate", methods=["POST"])
def api_generate():
    """Handle the backend's /api/generate format by forwarding to generate()"""
    return generate()

# Also add /api/analyze and /api/vision-chat for the backend
@app.route("/api/analyze", methods=["POST"])
def api_analyze():
    """Image analysis - return a placeholder since we don't have a vision model here"""
    data = request.get_json(force=True)
    return jsonify({
        "success": True,
        "description": "Image analysis is not available on this endpoint.",
        "model": "none",
    })

@app.route("/api/vision-chat", methods=["POST"])
def api_vision_chat():
    """Vision chat - return a placeholder"""
    data = request.get_json(force=True)
    return jsonify({
        "success": True,
        "answer": "Vision chat is not available on this endpoint.",
        "model": "none",
    })

if __name__ == "__main__":
    load_model()
    app.run(host="0.0.0.0", port=8000)