import os
import threading
import time

import gradio as gr
import numpy as np
import spaces
import torch
from huggingface_hub import hf_hub_download, snapshot_download
from omegaconf import OmegaConf
from PIL import Image
from transformers import StoppingCriteriaList

from minigpt4.conversation.conversation import (
    CONV_VISION_Vicuna0, Chat, StoppingCriteriaSub)
from minigpt4.models.minigpt4 import MiniGPT4
from minigpt4.processors.blip_processors import Blip2ImageEvalProcessor

# Vicuna v0 normally needs LLaMA-1 base weights plus a delta, and LLaMA-1 was never
# publicly distributed. These are the authors' own prepared weights, which is the
# only reason this model can still be assembled at all.
VICUNA_REPO = "Vision-CAIR/vicuna-7b"
# The trained linear projection. Everything else (ViT-G, Q-Former) stays frozen.
CKPT_REPO, CKPT_FILE = "ckpt/minigpt4-7B", "prerained_minigpt4_7b.pth"

MODEL = None
CHAT = None
VISION_OK = None
LOCK = threading.Lock()
_ON_GPU = False


def build():
    global MODEL, CHAT, VISION_OK
    vicuna = snapshot_download(VICUNA_REPO)
    ckpt = hf_hub_download(CKPT_REPO, CKPT_FILE)

    cfg = OmegaConf.load("minigpt4/configs/models/minigpt4_vicuna0.yaml").model
    cfg.llama_model = vicuna
    cfg.ckpt = ckpt
    cfg.low_resource = False          # 24 GB of A10G, no need for the 8-bit path
    cfg.prompt_template = "###Human: {} ###Assistant: "
    cfg.max_txt_len = 160
    cfg.end_sym = "###"

    MODEL = MiniGPT4.from_config(cfg).eval()
    vis = Blip2ImageEvalProcessor.from_config(
        OmegaConf.create({"name": "blip2_image_eval", "image_size": 224}))

    # "###" and "##" as stop sequences, straight from the reference demo.
    stops = [torch.tensor(ids) for ids in ([835], [2277, 29937])]
    CHAT = Chat(MODEL, vis, device="cpu",
                stopping_criteria=StoppingCriteriaList([StoppingCriteriaSub(stops=stops)]))
    print("[mg4] model built on CPU", flush=True)


def to_gpu():
    """ZeroGPU only exposes CUDA inside a @spaces.GPU call, so the move happens on
    first use rather than at import."""
    global _ON_GPU
    if _ON_GPU:
        return
    dev = torch.device("cuda")
    MODEL.to(dev)
    CHAT.device = dev
    CHAT.stopping_criteria = StoppingCriteriaList(
        [StoppingCriteriaSub(stops=[torch.tensor(i).to(dev) for i in ([835], [2277, 29937])])])
    _ON_GPU = True
    print("[mg4] moved to CUDA", flush=True)


def _embed(img):
    lst = []
    conv = CONV_VISION_Vicuna0.copy()
    CHAT.upload_img(img, conv, lst)
    CHAT.encode_img(lst)
    return lst[0].detach().float().cpu().numpy()


@spaces.GPU(duration=120)
def selftest():
    """The CPU ggml port returned bit-identical embeddings for any two images, which
    is what made every answer confabulated. Prove this build does not, and time the
    encode while we are here so the header can quote a real number."""
    to_gpu()
    black = Image.new("RGB", (448, 448), (0, 0, 0))
    patt = Image.fromarray(
        (np.indices((448, 448)).sum(0) % 256).astype("uint8")).convert("RGB")
    a = _embed(black)
    t = time.time()
    b = _embed(patt)
    enc = time.time() - t
    diff = float(np.abs(a - b).max())
    ok = diff > 1e-4
    print(f"[mg4] SELFTEST vision responds to input: {ok} (maxdiff={diff:.6f} "
          f"shape={a.shape} encode={enc:.2f}s)", flush=True)
    return ok, enc


def example_run(image, prompt):
    """The example row carries only the two inputs, so the sliders keep their
    defaults. Routed through chat() so a cached example is the same code path
    the Send button runs."""
    return chat(image, prompt, [], 300, 0.7, 0.9)


@spaces.GPU(duration=120)
def chat(image, prompt, history, max_new, temperature, top_p):
    history = history or []
    if image is None:
        raise gr.Error("Upload an image first. This model always answers about an image.")
    if not (prompt or "").strip():
        raise gr.Error("Type a question about the image.")
    if not LOCK.acquire(blocking=False):
        raise gr.Error("Another request is running. Try again in a moment.")
    try:
        to_gpu()
        t0 = time.time()
        conv = CONV_VISION_Vicuna0.copy()
        img_list = []
        CHAT.upload_img(image, conv, img_list)
        CHAT.encode_img(img_list)
        enc = time.time() - t0

        CHAT.ask(prompt, conv)
        t1 = time.time()
        answer = CHAT.answer(conv=conv, img_list=img_list,
                             num_beams=1, temperature=float(temperature),
                             top_p=float(top_p), max_new_tokens=int(max_new),
                             max_length=2000)[0]
        dt = time.time() - t1
        answer = answer.replace("<Img>", "").replace("</Img>", "").strip()
        note = f"\n\n`encode {enc:.1f}s · generate {dt:.1f}s`"
        return history + [{"role": "user", "content": prompt},
                          {"role": "assistant", "content": answer + note}], ""
    finally:
        LOCK.release()


CSS = """
footer{display:none}
.gap{gap:0}
.block{padding:2px}
#img .upload-container,#img .file-preview{max-height:70px}
#img img{max-height:150px;object-fit:contain}
#hdr{max-height:32px}
#hdr p{margin:0}
:where(button,input,textarea,select,[role="button"],[tabindex]):focus-visible{
  outline:2px solid #7c5cff!important;outline-offset:2px}
@media (prefers-reduced-motion:reduce){*,*::before,*::after{
  animation-duration:.01ms!important;transition-duration:.01ms!important}}
"""

build()
VISION_OK, ENC_S = selftest()

# One line, resolved backend and measured cost, so the wall clock is predictable
# before anyone uploads anything. Generation is ~9 ms/token on this hardware.
GPU_NAME = "ZeroGPU A10G"
# Sub-second encodes read as "0.0s" at one decimal, which looks like broken telemetry.
ENC_TXT = f"{ENC_S * 1000:.0f}ms" if ENC_S < 0.5 else f"{ENC_S:.1f}s"
HEADER = (f"**MiniGPT-4 7B** · uncensored captioning · Vicuna v0 · {GPU_NAME} · "
          f"~{ENC_TXT} encode + ~{300 * 0.009:.0f}s per 300 tokens"
          + ("" if VISION_OK else "  ·  ⚠ **vision self-test FAILED, answers are "
                                  "not grounded in the image**"))

with gr.Blocks(fill_width=True) as demo:
    gr.Markdown(HEADER, elem_id="hdr")
    with gr.Row():
        with gr.Column(scale=3):
            img = gr.Image(type="pil", show_label=False, elem_id="img", height=150)
            box = gr.Chatbot(show_label=False, height=230)
            with gr.Row():
                msg = gr.Textbox(placeholder="Ask about the image...", show_label=False,
                                 container=False, lines=2, max_lines=2, scale=5)
                go = gr.Button("Send", variant="primary", scale=1)
            wipe = gr.Button("New conversation", size="sm")
        with gr.Column(scale=1):
            gr.Markdown(
                "Vicuna v0 predates the refusal training that arrived later in 2023, so "
                "it describes what it is shown instead of declining. It also "
                "hallucinates freely. A 2023 model, not a current one."
            )
            with gr.Accordion("Settings", open=False):
                max_new = gr.Slider(32, 512, value=300, step=32,
                                    label="Max new tokens (300 ≈ 3s)")
                temperature = gr.Slider(0.1, 1.5, value=0.7, step=0.05,
                                        label="Temperature",
                                        info="Higher wanders further from the image.")
                top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top P")

    go.click(chat, [img, msg, box, max_new, temperature, top_p], [box, msg])
    msg.submit(chat, [img, msg, box, max_new, temperature, top_p], [box, msg])
    wipe.click(lambda: ([], ""), outputs=[box, msg], queue=False)

    # Exactly one row. Every row is a full pipeline run if caching is ever enabled.
    # With no fn the example only fills the inputs and the visitor still has to
    # press Send, which reads as a broken example. Give it the handler and cache
    # it. Lazy rather than eager because caching at startup would need the GPU
    # before ZeroGPU has handed one over.
    gr.Examples(examples=[["examples/cats.jpg", "Describe this image in detail."]],
                inputs=[img, msg], outputs=[box, msg],
                fn=example_run, cache_examples=True, cache_mode="lazy",
                label="Example")

demo.queue(default_concurrency_limit=1, max_size=8).launch(css=CSS, show_error=True)