"""Inference backend — one place that owns the models and GPU calls.

Everything model-related goes through `vision_generate()` and `text_generate()`
so the rest of the app doesn't care *where* inference runs. Today that's local
on the Space's ZeroGPU; the same two functions can later be backed by Modal
(set BB_INFERENCE=modal) without touching extract/categorize/agent/chat.

Models (both eligible for the hackathon):
- Vision: MiniCPM-V-4.6 (1.3B) — OCR/extraction.
- Text:  MiniCPM4.1-8B (8B reasoning) — understanding, categorising, the agent.
"""

from __future__ import annotations

import os
import re

VISION_MODEL_ID = "openbmb/MiniCPM-V-4.6"
# On-Space text model must be transformers-5.12-native (MiniCPM-V-4.6 forces
# transformers>=5.7). MiniCPM4.1-8B's trust_remote_code is written for ~4.56 and
# RuntimeErrors on 5.12, so the 8B can only run isolated on the Modal backend.
LOCAL_TEXT_MODEL_ID = "openbmb/MiniCPM5-1B"
MODAL_TEXT_MODEL_ID = "openbmb/MiniCPM4.1-8B"
BACKEND = os.environ.get("BB_INFERENCE", "local").lower()  # "local" | "modal"
TEXT_MODEL_ID = MODAL_TEXT_MODEL_ID if BACKEND == "modal" else LOCAL_TEXT_MODEL_ID

# Vision detail (small receipt text).
DOWNSAMPLE_MODE = "4x"
MAX_SLICE_NUMS = 36

_THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL | re.IGNORECASE)


# --------------------------------------------------------------------------- #
# ZeroGPU decorator (no-op locally)
# --------------------------------------------------------------------------- #
try:
    import spaces  # type: ignore

    gpu_decorator = spaces.GPU
except Exception:

    def gpu_decorator(func=None, **_kwargs):  # type: ignore
        if func is None:
            return lambda f: f
        return func


# --------------------------------------------------------------------------- #
# Lazy model loaders (cached)
# --------------------------------------------------------------------------- #
_vision = None  # (model, processor)
_text = None    # (model, tokenizer)


def _load_vision():
    global _vision
    if _vision is not None:
        return _vision
    from transformers import AutoModelForImageTextToText, AutoProcessor

    processor = AutoProcessor.from_pretrained(VISION_MODEL_ID)
    model = AutoModelForImageTextToText.from_pretrained(
        VISION_MODEL_ID, torch_dtype="auto", device_map="auto"
    ).eval()
    _vision = (model, processor)
    return _vision


def _shim_text_remote_code():
    """MiniCPM4.1-8B's trust_remote_code was written for transformers ~4.56 and
    imports a few internals removed in 5.x. Inject safe equivalents so it loads
    under the 5.7 we need for MiniCPM-V-4.6. (If runtime APIs also diverge, the
    real fix is Modal — running the 8B in its own transformers env.)"""
    try:
        import torch.nn as nn
        import transformers.pytorch_utils as pu
        if not hasattr(pu, "is_torch_greater_or_equal_than_1_13"):
            pu.is_torch_greater_or_equal_than_1_13 = True
        if not hasattr(pu, "ALL_LAYERNORM_LAYERS"):
            pu.ALL_LAYERNORM_LAYERS = [nn.LayerNorm]
        import transformers.utils.import_utils as iu
        if not hasattr(iu, "is_torch_fx_available"):
            iu.is_torch_fx_available = lambda: False
    except Exception as e:  # pragma: no cover
        print(f"[inference] text shim failed: {e}")


def _load_text():
    global _text
    if _text is not None:
        return _text
    _shim_text_remote_code()
    from transformers import AutoModelForCausalLM, AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_ID, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        TEXT_MODEL_ID, torch_dtype="auto", device_map="auto", trust_remote_code=True
    ).eval()
    _text = (model, tokenizer)
    return _text


def preload():
    """Load local models once at import. Vision always runs on the Space; the
    text model loads locally only when the text backend is local (with Modal it
    runs remotely, so we skip the heavy local load)."""
    try:
        _load_vision()
    except Exception as e:  # pragma: no cover
        print(f"[inference] vision deferred: {e}")
    if BACKEND == "local":
        try:
            _load_text()
        except Exception as e:  # pragma: no cover
            print(f"[inference] text deferred: {e}")


# --------------------------------------------------------------------------- #
# Generation
# --------------------------------------------------------------------------- #
@gpu_decorator
def _vision_local(image, system: str, user: str, max_new_tokens: int) -> str:
    model, processor = _load_vision()
    messages = [
        {"role": "system", "content": system},
        {"role": "user", "content": [{"type": "image", "image": image},
                                      {"type": "text", "text": user}]},
    ]
    inputs = processor.apply_chat_template(
        messages, tokenize=True, add_generation_prompt=True, return_dict=True,
        return_tensors="pt", downsample_mode=DOWNSAMPLE_MODE, max_slice_nums=MAX_SLICE_NUMS,
    ).to(model.device)
    out = model.generate(**inputs, downsample_mode=DOWNSAMPLE_MODE,
                         max_new_tokens=max_new_tokens, do_sample=False)
    trimmed = [o[len(i):] for i, o in zip(inputs["input_ids"], out)]
    return str(processor.batch_decode(trimmed, skip_special_tokens=True)[0])


@gpu_decorator
def _text_local(messages: list[dict], max_new_tokens: int, enable_thinking: bool) -> str:
    model, tokenizer = _load_text()
    try:
        inputs = tokenizer.apply_chat_template(
            messages, tokenize=True, add_generation_prompt=True,
            enable_thinking=enable_thinking, return_dict=True, return_tensors="pt",
        ).to(model.device)
    except TypeError:  # template doesn't accept enable_thinking
        inputs = tokenizer.apply_chat_template(
            messages, tokenize=True, add_generation_prompt=True,
            return_dict=True, return_tensors="pt",
        ).to(model.device)
    out = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
    text = tokenizer.decode(out[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
    return _THINK_RE.sub("", str(text)).strip()


def vision_generate(image, system: str, user: str, max_new_tokens: int = 1024) -> str:
    """Run the vision model on an image (always local — it needs transformers 5.x)."""
    return _vision_local(image, system, user, max_new_tokens)


def text_generate(messages: list[dict], max_new_tokens: int = 512,
                  enable_thinking: bool = False) -> str:
    """Run the text model on chat messages. Returns text (<think> stripped)."""
    if BACKEND == "modal":
        from core import modal_backend  # lazy
        return modal_backend.text_generate(messages, max_new_tokens, enable_thinking)
    return _text_local(messages, max_new_tokens, enable_thinking)


preload()