import os

# Must be set before torch/diffusers import
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

# MUST come before torch / diffusers (ZeroGPU scanner)
import spaces

import torch
import gradio as gr
import asyncio
import httpx
import tempfile
import traceback

from fastapi import Request
from fastapi.responses import JSONResponse
from starlette.routing import Route as StarletteRoute

# ✅ CORRECT pipeline for Juggernaut Z — NOT Lumina2Pipeline
from diffusers import ZImagePipeline


# ============================================================
# CONFIG
# ============================================================
MODEL_ID = os.getenv(
    "MODEL_ID",
    "RunDiffusion/Juggernaut-Z-Image",
)
TELEGRAM_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN", "")
TELEGRAM_API = f"https://api.telegram.org/bot{TELEGRAM_TOKEN}" if TELEGRAM_TOKEN else ""
MAX_IMAGES = int(os.getenv("MAX_IMAGES_PER_REQUEST", "10"))
GPU_DURATION = int(os.getenv("GPU_DURATION_SEC", "120"))

# Juggernaut Z recommended settings (official model card)
STEPS = int(os.getenv("GEN_STEPS", "35"))
GUIDANCE = float(os.getenv("GEN_GUIDANCE", "6.0"))
GEN_SIZE = int(os.getenv("GEN_SIZE", "1024"))


# ============================================================
# MODEL LOAD (module level — ZeroGPU streams weights on first GPU call)
# ============================================================
print(f"[BOOT] Loading Juggernaut Z pipeline: {MODEL_ID}")

# Use ZImagePipeline. If your diffusers version rejects it, fall back to:
# from diffusers import DiffusionPipeline
# pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
pipe = ZImagePipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
)

print("[BOOT] Pipeline loaded (weights on CPU/disk; ZeroGPU will stream to VRAM)")


# ============================================================
# TELEGRAM HELPERS
# ============================================================
async def send_message(chat_id: int, text: str) -> None:
    if not TELEGRAM_TOKEN:
        return
    try:
        async with httpx.AsyncClient(timeout=15.0) as client:
            await client.post(
                f"{TELEGRAM_API}/sendMessage",
                json={
                    "chat_id": chat_id,
                    "text": text,
                    "parse_mode": "HTML",
                    "disable_web_page_preview": True,
                },
            )
    except Exception as e:
        print(f"[TG] send_message failed: {e}")


async def send_photo(chat_id: int, path: str, caption: str = "") -> None:
    if not TELEGRAM_TOKEN:
        return
    try:
        async with httpx.AsyncClient(timeout=120.0) as client:
            with open(path, "rb") as f:
                await client.post(
                    f"{TELEGRAM_API}/sendPhoto",
                    data={"chat_id": chat_id, "caption": caption, "parse_mode": "HTML"},
                    files={"photo": f},
                )
    except Exception as e:
        print(f"[TG] send_photo failed: {e}")


# ============================================================
# GPU FUNCTION — the ONLY @spaces.GPU function in this Space
# ============================================================
@spaces.GPU(duration=GPU_DURATION)
def generate_image_gpu(prompt: str, seed: int = -1) -> str:
    """
    Generate one hyper-realistic image with Juggernaut Z.
    Official recommended settings: 35 steps, CFG 6.0.
    """
    # Move to GPU on first call. ZeroGPU caches weights in VRAM afterwards.
    if pipe.device.type != "cuda":
        pipe.to("cuda")

    generator = None
    if seed != -1:
        generator = torch.Generator(device="cuda").manual_seed(int(seed))

    result = pipe(
        prompt=prompt,
        generator=generator,
        num_inference_steps=STEPS,
        guidance_scale=GUIDANCE,
        height=GEN_SIZE,
        width=GEN_SIZE,
    )

    image = result.images[0]

    tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
    image.save(tmp)
    return tmp


# ============================================================
# TASK QUEUE + WORKER
# ============================================================
task_queue: asyncio.Queue = asyncio.Queue()
_worker_task: asyncio.Task | None = None


async def worker() -> None:
    print("[WORKER] started")
    while True:
        chat_id = None
        try:
            chat_id, prompt, num = await task_queue.get()
            await send_message(chat_id, f"🎨 Generating {num} hyper-realistic image(s)...")
            for i in range(num):
                await send_message(chat_id, f"⏳ Image {i + 1}/{num} — rendering (35 steps)...")
                try:
                    # Run the sync GPU function in a thread so we don't block the loop
                    path = await asyncio.to_thread(generate_image_gpu, prompt)
                    await send_photo(chat_id, path, f"✅ Image {i + 1}/{num}")
                    try:
                        os.unlink(path)
                    except OSError:
                        pass
                except Exception as e:
                    print(f"[WORKER] image {i + 1} failed: {e}\n{traceback.format_exc()}")
                    await send_message(chat_id, f"❌ Image {i + 1} failed: {e}")
            await send_message(chat_id, "🎉 All done!")
        except Exception as e:
            print(f"[WORKER] fatal: {e}\n{traceback.format_exc()}")
            if chat_id:
                try:
                    await send_message(chat_id, f"❌ Worker error: {e}")
                except Exception:
                    pass
        finally:
            task_queue.task_done()


# ============================================================
# FASTAPI ROUTES
# ============================================================
async def telegram_webhook(request: Request):
    global _worker_task
    # Lazy-start the worker inside the running event loop
    if _worker_task is None:
        _worker_task = asyncio.create_task(worker())

    try:
        update = await request.json()
        msg = update.get("message", {})
        chat_id = msg.get("chat", {}).get("id")
        text = msg.get("text", "")

        if chat_id and text:
            if text.startswith("/start"):
                await send_message(
                    chat_id,
                    "👋 <b>Hyper-Realistic Image Bot (Juggernaut Z)</b>\n\n"
                    f"Send <code>/imagine N &lt;prompt&gt;</code> (N = 1..{MAX_IMAGES}).\n"
                    "Example: <code>/imagine 2 cinematic portrait of an old fisherman, dramatic lighting</code>",
                )
            elif text.startswith("/imagine"):
                parts = text.split(maxsplit=2)
                if len(parts) >= 3:
                    try:
                        num = int(parts[1])
                        if not 1 <= num <= MAX_IMAGES:
                            raise ValueError
                    except ValueError:
                        await send_message(chat_id, f"❌ Number must be 1..{MAX_IMAGES}")
                        return JSONResponse({"ok": True})
                    prompt = parts[2]
                    await task_queue.put((chat_id, prompt, num))
                    await send_message(chat_id, f"✅ Queued {num} image(s): <i>{prompt}</i>")
                else:
                    await send_message(chat_id, "Usage: <code>/imagine N prompt</code>")
    except Exception as e:
        print(f"[WEBHOOK] error: {e}\n{traceback.format_exc()}")

    return JSONResponse({"ok": True})


async def health_check(request: Request):
    return JSONResponse({
        "ok": True,
        "status": "running",
        "queue_size": task_queue.qsize(),
        "worker_alive": _worker_task is not None and not _worker_task.done(),
    })


# ============================================================
# GRADIO DEMO — hidden button satisfies the ZeroGPU scanner
# ============================================================
with gr.Blocks() as demo:
    gr.Markdown("# 🎨 Hyper-Realistic Image Bot (Juggernaut Z)\n\nServes a Telegram bot via `/telegram`.")
    _btn = gr.Button("init-gpu", visible=False)
    _in = gr.Textbox(value="cinematic portrait, dramatic lighting, photorealistic", visible=False)
    _out = gr.File(visible=False)
    _btn.click(fn=generate_image_gpu, inputs=[_in], outputs=[_out])


# ============================================================
# LAUNCH — proven pattern from your history
# ============================================================
demo.launch(
    server_name="0.0.0.0",
    server_port=7860,
    share=False,
    ssr_mode=False,             # prevents Node proxy from swallowing /telegram
    prevent_thread_lock=True,   # returns control so we can inject routes
)

# Inject routes at index 0 so they beat Gradio's catch-all Mount
live_app = demo.app
live_app.router.routes.insert(
    0, StarletteRoute("/telegram", telegram_webhook, methods=["POST"])
)
live_app.router.routes.insert(
    0, StarletteRoute("/health", health_check, methods=["GET"])
)
live_app.build_middleware_stack()

print("[BOOT] Routes injected: POST /telegram, GET /health")
print("[BOOT] Set webhook → https://<space>.hf.space/telegram")

demo.block_thread()