"""
OpenAI-compatible vision shim for an uncensored Gemma-3 12B judge.

Loads llmfan46/gemma-3-12b-it-ultra-uncensored-heretic in 4-bit (bitsandbytes)
and exposes POST /v1/chat/completions accepting text + image_url content, so the
FLUX pipeline's GemmaClient can talk to it unchanged.

Runs on a T4 (16 GB): 12B in nf4 ≈ 6-7 GB + vision tower + KV → comfortable.
"""
import base64
import io
import os
import time
import urllib.request

import torch
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from PIL import Image

MODEL_ID = os.environ.get("MODEL_ID", "llmfan46/gemma-4-E4B-it-ultra-uncensored-heretic")

app = FastAPI()
_proc = None
_model = None
_load_error = None


def _load():
    """Lazy-load the model on first request (keeps startup/health fast)."""
    global _proc, _model, _load_error
    if _model is not None or _load_error is not None:
        return
    try:
        # model-agnostic loader → works for gemma3 AND gemma4 (needs transformers>=5.13)
        from transformers import AutoModelForImageTextToText, AutoProcessor, BitsAndBytesConfig
        _proc = AutoProcessor.from_pretrained(MODEL_ID)
        bnb = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.float16,   # T4 has no bf16
            bnb_4bit_use_double_quant=True,
        )
        _model = AutoModelForImageTextToText.from_pretrained(
            MODEL_ID, quantization_config=bnb, device_map="auto", torch_dtype=torch.float16,
        )
        _model.eval()
    except Exception as e:  # surfaced via /health and request errors
        _load_error = f"{type(e).__name__}: {e}"


def _load_image(url: str) -> Image.Image:
    if url.startswith("data:"):
        b64 = url.split(",", 1)[1]
        return Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
    with urllib.request.urlopen(url, timeout=30) as r:
        return Image.open(io.BytesIO(r.read())).convert("RGB")


def _to_gemma_messages(messages):
    """Convert OpenAI messages (string or [parts]) into the transformers
    multimodal chat format, inlining PIL images."""
    conv = []
    for m in messages:
        role, content = m.get("role", "user"), m.get("content", "")
        if isinstance(content, str):
            conv.append({"role": role, "content": [{"type": "text", "text": content}]})
            continue
        parts = []
        for p in content:
            if p.get("type") == "text":
                parts.append({"type": "text", "text": p.get("text", "")})
            elif p.get("type") == "image_url":
                parts.append({"type": "image", "image": _load_image(p["image_url"]["url"])})
        conv.append({"role": role, "content": parts})
    return conv


@app.get("/health")
def health():
    return {"status": "ok", "model_loaded": _model is not None, "load_error": _load_error}


@app.post("/v1/chat/completions")
async def chat(req: Request):
    _load()
    if _load_error:
        return JSONResponse({"error": _load_error}, status_code=500)
    body = await req.json()
    try:
        conv = _to_gemma_messages(body.get("messages", []))
        inputs = _proc.apply_chat_template(
            conv, add_generation_prompt=True, tokenize=True,
            return_dict=True, return_tensors="pt",
        ).to(_model.device)
        in_len = inputs["input_ids"].shape[1]
        max_new = int(body.get("max_tokens", 512))
        temp = float(body.get("temperature", 0.7))
        with torch.inference_mode():
            out = _model.generate(
                **inputs, max_new_tokens=max_new,
                do_sample=temp > 0, temperature=max(temp, 0.01),
            )
        gen = out[0][in_len:]
        text = _proc.decode(gen, skip_special_tokens=True).strip()
        return {
            "id": "chatcmpl-gemma", "object": "chat.completion",
            "created": int(time.time()), "model": MODEL_ID,
            "choices": [{"index": 0, "message": {"role": "assistant", "content": text},
                         "finish_reason": "stop"}],
            "usage": {"prompt_tokens": int(in_len), "completion_tokens": int(gen.shape[0]),
                      "total_tokens": int(in_len + gen.shape[0])},
        }
    except Exception as e:
        return JSONResponse({"error": f"{type(e).__name__}: {e}"}, status_code=500)


@app.get("/")
def root():
    return {"service": "gemma-vision-judge", "model": MODEL_ID,
            "endpoints": ["/health", "/v1/chat/completions"]}