import atexit
import gc
import os
import re
import shutil
import subprocess
import threading
from dataclasses import dataclass

import gradio as gr
from huggingface_hub import hf_hub_download

import litert_lm


litert_lm.set_min_log_severity(litert_lm.LogSeverity.ERROR)

MODEL_CACHE_DIR = os.environ.get("HF_HOME", "/tmp/huggingface")
LITERT_CACHE_DIR = os.environ.get("LITERT_LM_CACHE_DIR", "/tmp/litert-lm-cache")
MAX_OUTPUT_TOKENS = int(os.environ.get("MAX_OUTPUT_TOKENS", "220"))
FORCE_CPU = os.environ.get("FORCE_CPU", "").strip().lower() in {"1", "true", "yes"}


@dataclass(frozen=True)
class ModelSpec:
    repo_id: str
    filename: str


MODELS = {
    "Gemma 4 E2B IT (LiteRT-LM)": ModelSpec(
        repo_id="litert-community/gemma-4-E2B-it-litert-lm",
        filename="gemma-4-E2B-it.litertlm",
    ),
    "Gemma 4 E4B IT (LiteRT-LM)": ModelSpec(
        repo_id="litert-community/gemma-4-E4B-it-litert-lm",
        filename="gemma-4-E4B-it.litertlm",
    ),
}

SYSTEM_INSTRUCTION = """You are a precise image-description engine.
Describe only what is visible in the image. Do not invent context, identities, emotions, brands, locations, dates, or text that is not legible.
Write a compact description in one to four complete sentences, always fewer than five sentences.
Include the main subject, setting, notable objects, actions, composition, lighting, style, and any clearly readable text when relevant.
Return only the description. Do not include a title, bullets, markdown, analysis, caveats, or phrases such as "The image shows"."""

USER_PROMPT = """Inspect the image carefully and produce a sufficient visual description only.
Be specific but concise. Mention uncertainty only when needed for ambiguous visible details.
Keep the answer under five sentences."""

_model_lock = threading.RLock()
_engines = {}
_engine_contexts = {}
_engine_backends = {}
_preload_started = False


def _download_model(model_name: str) -> str:
    spec = MODELS[model_name]
    return hf_hub_download(
        repo_id=spec.repo_id,
        filename=spec.filename,
        cache_dir=MODEL_CACHE_DIR,
    )


def _has_nvidia_gpu() -> bool:
    if FORCE_CPU:
        return False

    visible_devices = os.environ.get("NVIDIA_VISIBLE_DEVICES", "").strip().lower()
    if visible_devices and visible_devices not in {"none", "void"}:
        return True

    if os.path.isdir("/proc/driver/nvidia/gpus"):
        return True

    if shutil.which("nvidia-smi") is None:
        return False

    try:
        result = subprocess.run(
            ["nvidia-smi", "-L"],
            check=False,
            capture_output=True,
            text=True,
            timeout=3,
        )
    except (OSError, subprocess.SubprocessError):
        return False
    return result.returncode == 0 and "GPU" in result.stdout


def _backend(name: str):
    return getattr(litert_lm.Backend, name, None)


def _preferred_backends():
    gpu_backend = _backend("GPU")
    cpu_backend = _backend("CPU")
    if _has_nvidia_gpu() and gpu_backend is not None:
        return [("GPU", gpu_backend), ("CPU", cpu_backend)]
    return [("CPU", cpu_backend)]


def _close_engine(model_name: str) -> None:
    context = _engine_contexts.pop(model_name, None)
    engine = _engines.pop(model_name, None)
    _engine_backends.pop(model_name, None)

    if context is not None:
        context.__exit__(None, None, None)
    elif engine is not None:
        close = getattr(engine, "close", None)
        if callable(close):
            close()
    gc.collect()


def _close_all_engines() -> None:
    with _model_lock:
        for model_name in list(_engines):
            _close_engine(model_name)


def _get_engine(model_name: str):
    with _model_lock:
        if model_name in _engines:
            return _engines[model_name]

        model_path = _download_model(model_name)
        model_cache_dir = os.path.join(
            LITERT_CACHE_DIR,
            re.sub(r"[^A-Za-z0-9_.-]+", "_", model_name),
        )
        os.makedirs(model_cache_dir, exist_ok=True)

        last_error = None
        for backend_name, backend in _preferred_backends():
            if backend is None:
                continue
            try:
                engine_context = litert_lm.Engine(
                    model_path,
                    backend=backend,
                    vision_backend=backend,
                    cache_dir=model_cache_dir,
                )
                engine = engine_context.__enter__()
            except Exception as exc:
                last_error = exc
                print(
                    f"Failed to load {model_name} with {backend_name}; "
                    "trying the next backend.",
                    flush=True,
                )
                continue

            _engine_contexts[model_name] = engine_context
            _engines[model_name] = engine
            _engine_backends[model_name] = backend_name
            print(f"Loaded {model_name} with LiteRT-LM {backend_name} backend.", flush=True)
            return engine

        raise RuntimeError(f"Could not load {model_name}.") from last_error


def _preload_models() -> None:
    global _preload_started
    with _model_lock:
        if _preload_started:
            return
        _preload_started = True

    for model_name in MODELS:
        _get_engine(model_name)


def _extract_text(response: dict) -> str:
    parts = []
    for item in response.get("content", []):
        if item.get("type") == "text":
            parts.append(item.get("text", ""))
    return "".join(parts).strip()


def _trim_to_four_sentences(text: str) -> str:
    text = re.sub(r"\s+", " ", text).strip()
    text = re.sub(r"^(?:description|answer)\s*:\s*", "", text, flags=re.IGNORECASE)
    sentences = re.findall(r"[^.!?]+[.!?]+(?:\s|$)|[^.!?]+$", text)
    return " ".join(sentence.strip() for sentence in sentences[:4]).strip()


def describe_image(
    image_path: str,
    model_name: str,
    system_prompt: str,
    user_prompt: str,
    progress=gr.Progress(track_tqdm=False),
) -> str:
    if not image_path:
        raise gr.Error("Upload an image first.")

    if model_name not in MODELS:
        raise gr.Error("Choose a supported Gemma 4 LiteRT-LM model.")

    system_prompt = (system_prompt or SYSTEM_INSTRUCTION).strip()
    user_prompt = (user_prompt or USER_PROMPT).strip()

    progress(0.1, desc="Preparing model")
    engine = _get_engine(model_name)

    user_message = {
        "role": "user",
        "content": [
            {"type": "image", "path": image_path},
            {"type": "text", "text": user_prompt},
        ],
    }

    generation_kwargs = {"max_tokens": MAX_OUTPUT_TOKENS}

    progress(0.6, desc="Describing image")
    with _model_lock:
        with engine.create_conversation(
            messages=[
                {
                    "role": "system",
                    "content": [{"type": "text", "text": system_prompt}],
                }
            ],
        ) as conversation:
            try:
                response = conversation.send_message(user_message, **generation_kwargs)
            except TypeError:
                response = conversation.send_message(user_message)

    description = _trim_to_four_sentences(_extract_text(response))
    if not description:
        raise gr.Error("The model did not return a description. Try a clearer image or the other model.")
    return description


atexit.register(_close_all_engines)

print(
    "Starting Gemma 4 Image Description. "
    f"GPU visible: {_has_nvidia_gpu()}; force CPU: {FORCE_CPU}.",
    flush=True,
)
_preload_models()


theme = gr.themes.Soft(
    primary_hue="blue",
    neutral_hue="slate",
    radius_size="sm",
    text_size="md",
)

with gr.Blocks(title="Gemma 4 Image Description") as demo:
    gr.Markdown("# Gemma 4 Image Description")
    gr.Markdown("Upload an image and choose a LiteRT-LM model. Models are preloaded at startup and use GPU when available.")

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(
                label="Image",
                type="filepath",
                sources=["upload", "clipboard"],
                height=420,
            )
            model_input = gr.Radio(
                choices=list(MODELS.keys()),
                value="Gemma 4 E2B IT (LiteRT-LM)",
                label="Model",
            )
            with gr.Accordion("Prompts", open=False):
                system_prompt_input = gr.Textbox(
                    label="System prompt",
                    value=SYSTEM_INSTRUCTION,
                    lines=7,
                )
                user_prompt_input = gr.Textbox(
                    label="User prompt",
                    value=USER_PROMPT,
                    lines=4,
                )
            submit = gr.Button("Describe Image", variant="primary")

        with gr.Column(scale=1):
            output = gr.Textbox(
                label="Description",
                lines=8,
            )

    submit.click(
        fn=describe_image,
        inputs=[image_input, model_input, system_prompt_input, user_prompt_input],
        outputs=output,
        api_name="describe",
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(theme=theme, ssr_mode=False)