"""
SafeBite — allergen / ingredient label scanner (HF Build Small Hackathon · Backyard AI track)

Snap a product's ingredient label, pick your allergens/diet, and SafeBite reads the label
on-device (MiniCPM-V 4.6, ~1.3B, Apache-2.0) and flags what to avoid — works in the grocery
aisle with no signal, and your health profile never leaves the device.

Pipeline (each step logs its I/O to a trace):
  extract (vision model call)  ->  normalize (rules)  ->  match (alias dictionary, rules)
  ->  advise (rules)  ->  assemble
Only the extract step calls the model: a ~1.3B model reads labels well (OCR) but reasons
poorly, so allergen matching is a deterministic alias dictionary — safer and transparent.

No cloud LLM APIs. Inference path verified against the official openbmb/MiniCPM-V-4.6-Demo:
AutoProcessor + MiniCPMV4_6ForConditionalGeneration + apply_chat_template + generate.
"""

from __future__ import annotations

import json
import re
import tempfile
import time
from dataclasses import dataclass, field
from typing import Any

import gradio as gr
import spaces
import torch
from PIL import Image
from transformers import AutoProcessor, MiniCPMV4_6ForConditionalGeneration

# --------------------------------------------------------------------------- #
# Model loading (module level so ZeroGPU keeps it warm across calls)
# --------------------------------------------------------------------------- #
MODEL_ID = "openbmb/MiniCPM-V-4.6"
GPU_DURATION = 60

print(f"[safebite] loading processor: {MODEL_ID}", flush=True)
processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
print(f"[safebite] loading model: {MODEL_ID}", flush=True)
model = MiniCPMV4_6ForConditionalGeneration.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
    trust_remote_code=True,
    device_map="cuda",
).eval()


# --------------------------------------------------------------------------- #
# Core inference helper (verified signature)
# --------------------------------------------------------------------------- #
def _run_model(content: list[dict], max_new_tokens: int = 512) -> str:
    messages = [{"role": "user", "content": content}]
    has_image = any(item.get("type") == "image" for item in content)
    inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
        enable_thinking=False,
        processor_kwargs={
            "downsample_mode": "16x",
            "max_slice_nums": 9 if has_image else 1,
            "use_image_id": has_image,
        },
    ).to(model.device)
    for key, value in inputs.items():
        if isinstance(value, torch.Tensor) and torch.is_floating_point(value):
            inputs[key] = value.to(dtype=torch.bfloat16)
    generated = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False, downsample_mode="16x")
    trimmed = [out[len(inp):] for inp, out in zip(inputs.input_ids, generated)]
    return processor.batch_decode(trimmed, skip_special_tokens=True)[0].strip()


def _parse_json(raw: str) -> dict | None:
    """First balanced JSON object, tolerating trailing junk / stray braces."""
    start = raw.find("{")
    if start == -1:
        return None
    depth, in_str, esc = 0, False, False
    for i in range(start, len(raw)):
        ch = raw[i]
        if in_str:
            if esc:
                esc = False
            elif ch == "\\":
                esc = True
            elif ch == '"':
                in_str = False
        elif ch == '"':
            in_str = True
        elif ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                return _loads_lenient(raw[start:i + 1])
    return None


def _loads_lenient(s: str) -> dict | None:
    try:
        return json.loads(s)
    except json.JSONDecodeError:
        pass
    # Common model glitch: a stray '.' or ';' instead of ',' between a value and the next key.
    fixed = re.sub(r'([\]\}"])\s*[.;]\s*(\n\s*"[^"]+"\s*:)', r"\1,\2", s)
    try:
        return json.loads(fixed)
    except json.JSONDecodeError:
        return None


@dataclass
class Trace:
    steps: list[dict] = field(default_factory=list)

    def add(self, name: str, inp: Any, out: Any, ms: float) -> None:
        self.steps.append({"step": name, "input": inp, "output": out, "ms": round(ms, 1)})

    @property
    def total_ms(self) -> float:
        return sum(s["ms"] for s in self.steps)


# --------------------------------------------------------------------------- #
# Allergen / diet alias dictionary (the core asset — deterministic, transparent)
# Word-boundary matched. Over-flagging is safer than missing for allergies; the
# matched word + source are always shown so the user can judge each hit.
# --------------------------------------------------------------------------- #
ALIASES: dict[str, list[str]] = {
    "Dairy": ["dairy", "milk", "buttermilk", "butterfat", "milk solids", "milkfat", "casein", "caseinate",
              "caseinates", "sodium caseinate", "whey", "lactose", "ghee", "curd", "cream",
              "yogurt", "yoghurt", "cheese", "custard"],
    "Gluten": ["gluten", "wheat", "barley", "rye", "malt", "malted barley", "semolina", "spelt", "farro",
               "triticale", "durum", "couscous", "bulgur", "seitan", "graham"],
    "Tree nuts": ["tree nut", "tree nuts", "almond", "cashew", "walnut", "pecan", "hazelnut", "pistachio",
                  "macadamia", "brazil nut", "pine nut", "praline", "nut butter", "marzipan"],
    "Peanut": ["peanut", "peanuts", "groundnut", "arachis", "peanut butter"],
    "Egg": ["egg", "eggs", "albumin", "albumen", "ovalbumin", "globulin", "lysozyme", "meringue"],
    "Soy": ["soy", "soya", "soja", "edamame", "soybean", "tofu", "tempeh", "soy lecithin", "miso"],
    "Shellfish": ["shellfish", "crustacean", "crustaceans", "shrimp", "prawn", "crab", "lobster",
                  "crayfish", "crawfish", "scampi", "krill", "mollusk", "mollusc"],
    "Fish": ["fish", "anchovy", "anchovies", "cod", "tuna", "salmon", "haddock", "tilapia",
             "sardine", "fish sauce", "surimi"],
    "Sesame": ["sesame", "tahini", "benne", "sesamol", "gingelly"],
    "Sulfites": ["sulfite", "sulphite", "sulfites", "sulphites", "sodium bisulfite",
                 "potassium metabisulfite", "sulfur dioxide"],
    "Vegan": ["milk", "buttermilk", "casein", "caseinate", "whey", "lactose", "ghee", "cream",
              "yogurt", "cheese", "egg", "albumin", "honey", "gelatin", "gelatine", "carmine",
              "cochineal", "shellac", "rennet", "lard", "tallow", "isinglass", "beeswax",
              "fish", "anchovy", "meat", "chicken", "beef", "pork"],
    "Vegetarian": ["gelatin", "gelatine", "rennet", "lard", "tallow", "fish", "anchovy",
                   "fish sauce", "meat", "chicken", "beef", "pork", "isinglass", "carmine",
                   "cochineal"],
}
PROFILE_OPTIONS = list(ALIASES.keys())
_COMPILED = {cat: re.compile(r"\b(" + "|".join(re.escape(a) for a in sorted(aliases, key=len, reverse=True)) + r")\b", re.I)
             for cat, aliases in ALIASES.items()}


# --------------------------------------------------------------------------- #
# STEP 1 — EXTRACT (vision model call)
# --------------------------------------------------------------------------- #
EXTRACT_PROMPT = """You are reading a packaged-food label photo. Transcribe it as STRICT JSON, keys exactly:
{"product_name": str|null, "ingredients": [str, ...], "allergen_statements": [str, ...]}
Rules:
- "ingredients" = the items from the INGREDIENTS list, each as a SEPARATE string, split on commas, in order. Keep parenthetical sub-ingredients with their parent (e.g. "chocolate chips (sugar, cocoa, milk)").
- "allergen_statements" = each separate allergen summary line, copied VERBATIM, e.g. "Contains: Milk, Soy", "May contain tree nuts", "Made on shared equipment with peanuts". [] if there are none.
- Copy the words exactly as printed; do not add or infer allergens that are not written.
Return ONLY the JSON object, no markdown, no commentary."""


def step_extract(image: Image.Image, trace: Trace) -> dict | None:
    start = time.time()
    content = [{"type": "image", "image": image.convert("RGB")},
               {"type": "text", "text": EXTRACT_PROMPT}]
    raw = _run_model(content, max_new_tokens=512)
    parsed = _parse_json(raw)
    trace.add("extract", {"has_image": True}, parsed if parsed is not None else {"_raw": raw},
              (time.time() - start) * 1000)
    return parsed


# --------------------------------------------------------------------------- #
# STEP 2 — NORMALIZE (rules)
# --------------------------------------------------------------------------- #
def _as_list(value: Any) -> list[str]:
    if isinstance(value, list):
        return [str(v).strip() for v in value if str(v).strip()]
    if isinstance(value, str) and value.strip():
        return [p.strip() for p in re.split(r"[;,\n]", value) if p.strip()]
    return []


# Cross-contact ("may contain") triggers — classification done HERE, not by the model,
# because a ~1.3B model can't reliably tell "Contains" from "May contain".
_MAY_RE = re.compile(r"may contain|may be present|may include|traces? of|"
                     r"made (on|in)|produced (on|in)|packed (on|in)|processed (on|in)|"
                     r"shared (equipment|facility|line)|same (equipment|facility|line)", re.I)


def step_normalize(parsed: dict, trace: Trace) -> dict:
    start = time.time()
    statements = _as_list(parsed.get("allergen_statements"))
    # Back-compat: if a model still emits contains/may_contain keys, fold them in.
    statements += _as_list(parsed.get("contains")) + [f"may contain {x}" for x in _as_list(parsed.get("may_contain"))]
    contains, may_contain = [], []
    for s in statements:
        (may_contain if _MAY_RE.search(s) else contains).append(s)
    clean = {
        "product_name": (parsed.get("product_name") or "").strip() or None,
        "ingredients": _as_list(parsed.get("ingredients")),
        "contains": contains,
        "may_contain": may_contain,
        "statements": statements,
    }
    trace.add("normalize", parsed, clean, (time.time() - start) * 1000)
    return clean


# --------------------------------------------------------------------------- #
# STEP 3 — MATCH (deterministic alias dictionary)
# --------------------------------------------------------------------------- #
def step_match(facts: dict, profile: list[str], trace: Trace) -> list[dict]:
    start = time.time()
    buckets = {
        "ingredients": " , ".join(facts["ingredients"]),
        "contains": " , ".join(facts["contains"]),
        "may_contain": " , ".join(facts["may_contain"]),
    }
    hits: list[dict] = []
    seen = set()
    for cat in profile:
        rx = _COMPILED.get(cat)
        if not rx:
            continue
        for source, text in buckets.items():
            for matched in {m.lower() for m in rx.findall(text)}:
                key = (cat, matched, source)
                if key in seen:
                    continue
                seen.add(key)
                hits.append({"allergen": cat, "matched": matched, "source": source})
    trace.add("match", {"profile": profile, "buckets": {k: bool(v) for k, v in buckets.items()}},
              hits, (time.time() - start) * 1000)
    return hits


# --------------------------------------------------------------------------- #
# STEP 4 — ADVISE (rules: verdict tier)
# --------------------------------------------------------------------------- #
def step_advise(hits: list[dict], trace: Trace) -> dict:
    start = time.time()
    direct = [h for h in hits if h["source"] in ("ingredients", "contains")]
    trace_only = [h for h in hits if h["source"] == "may_contain"]
    if direct:
        tier, banner = "avoid", "🔴 AVOID"
    elif trace_only:
        tier, banner = "caution", "🟠 CAUTION"
    else:
        tier, banner = "clear", "🟢 No flagged allergens found"
    out = {"tier": tier, "banner": banner}
    trace.add("advise", {"direct": len(direct), "may_contain": len(trace_only)}, out, (time.time() - start) * 1000)
    return out


# --------------------------------------------------------------------------- #
# STEP 5 — ASSEMBLE
# --------------------------------------------------------------------------- #
DISCLAIMER = (
    "_Not medical advice and not a substitute for reading the physical label. AI reads labels "
    "from a single photo on a ~1.3B on-device model and can miss or misread text. If you have a "
    "severe allergy, always verify on the package itself._"
)
SOURCE_LABEL = {"ingredients": "Ingredients list", "contains": "“Contains” statement", "may_contain": "“May contain” / facility"}


def _hits_table(hits: list[dict]) -> str:
    rows = ["| Allergen / diet | Matched word | Found in |", "| --- | --- | --- |"]
    order = {"ingredients": 0, "contains": 1, "may_contain": 2}
    for h in sorted(hits, key=lambda x: order.get(x["source"], 3)):
        rows.append(f"| {h['allergen']} | `{h['matched']}` | {SOURCE_LABEL.get(h['source'], h['source'])} |")
    return "\n".join(rows)


def _render(facts: dict, hits: list[dict], advice: dict, profile: list[str], trace: Trace) -> str:
    name = facts.get("product_name") or "this product"
    parts = ["## 🥫 SafeBite — verdict", f"### {advice['banner']}"]

    if advice["tier"] == "avoid":
        parts.append(f"**{name}** contains something on your list — don't eat it without checking the package.")
    elif advice["tier"] == "caution":
        parts.append(f"**{name}** has a *may-contain* / shared-facility warning for your allergens — risky if you're sensitive.")
    else:
        parts.append(f"No items from your selected list were found in **{name}**. Still verify on the package.")

    if hits:
        parts.append("### 🚩 Flagged\n" + _hits_table(hits))

    if facts["ingredients"]:
        shown = ", ".join(facts["ingredients"][:40])
        parts.append("### 📋 Ingredients read\n" + shown)
    def _strip_label(s: str) -> str:
        return re.sub(r"^(contains|may contain|may also contain|may be present)\s*:?\s*", "", s, flags=re.I).strip()

    for label, key in (("Contains", "contains"), ("May contain", "may_contain")):
        if facts[key]:
            cleaned = [c for c in (_strip_label(s) for s in facts[key]) if c]
            if cleaned:
                parts.append(f"**{label}:** {', '.join(cleaned)}")

    parts.append(f"<sub>checked for: {', '.join(profile)} · pipeline: {' → '.join(s['step'] for s in trace.steps)} · {trace.total_ms:.0f} ms</sub>")
    parts.append(DISCLAIMER)
    return "\n\n".join(parts)


# --------------------------------------------------------------------------- #
# Orchestration
# --------------------------------------------------------------------------- #
def analyze(image: Image.Image | None, profile: list[str] | None) -> tuple[str, list]:
    if image is None:
        return "⚠️ Upload or snap a photo of the product's ingredient label to get started.", []
    if not profile:
        return "⚠️ Pick at least one allergen or diet to check against, then scan again.", []

    trace = Trace()
    parsed = step_extract(image, trace)
    if parsed is None or not _as_list(parsed.get("ingredients")) and not _as_list(parsed.get("allergen_statements")):
        raw = trace.steps[-1]["output"].get("_raw", "") if parsed is None else ""
        body = f"\n\nRaw model output:\n\n```\n{raw}\n```" if raw else ""
        return ("## 🥫 SafeBite\n\nI couldn't read an ingredients list in that photo. Try a sharper, "
                "well-lit close-up of the label." + body + "\n\n" + DISCLAIMER), trace.steps

    facts = step_normalize(parsed, trace)
    hits = step_match(facts, profile, trace)
    advice = step_advise(hits, trace)
    return _render(facts, hits, advice, profile, trace), trace.steps


def _trace_file(steps: list) -> str | None:
    """Write the run's pipeline trace to a downloadable JSON (Open Trace)."""
    if not steps:
        return None
    payload = {"app": "safebite", "model": MODEL_ID, "pipeline": [s["step"] for s in steps], "steps": steps}
    handle = tempfile.NamedTemporaryFile(mode="w", suffix="_safebite-trace.json", delete=False, encoding="utf-8")
    json.dump(payload, handle, indent=2, ensure_ascii=False)
    handle.close()
    return handle.name


@spaces.GPU(duration=GPU_DURATION)
def analyze_gpu(image: Image.Image | None, profile: list[str] | None) -> tuple[str, str | None]:
    markdown, steps = analyze(image, profile)
    return markdown, _trace_file(steps)


# --------------------------------------------------------------------------- #
# UI
# --------------------------------------------------------------------------- #
EXAMPLES = [
    ["assets/examples/label_granola.png", ["Peanut", "Dairy"]],
    ["assets/examples/label_cookie.png", ["Vegan"]],
    ["assets/examples/label_crackers.png", ["Tree nuts"]],
]


def build_ui() -> gr.Blocks:
    with gr.Blocks(title="SafeBite", theme=gr.themes.Soft()) as demo:
        gr.Markdown(
            "# 🥫 SafeBite\n"
            "Snap a product's **ingredient label**, pick what you avoid, and SafeBite reads it "
            "**on-device** on MiniCPM-V 4.6 and flags what to skip. Works in the aisle with no "
            "signal — and your health profile never leaves the device. _No cloud APIs._"
        )
        with gr.Row():
            with gr.Column(scale=1):
                image_in = gr.Image(type="pil", label="Ingredient label photo", height=340)
                profile_in = gr.CheckboxGroup(PROFILE_OPTIONS, label="I need to avoid…", value=["Peanut", "Dairy"])
                run_btn = gr.Button("Check label", variant="primary")
                gr.Examples(examples=EXAMPLES, inputs=[image_in, profile_in], label="Try an example")
            with gr.Column(scale=1):
                verdict = gr.Markdown("Your verdict will appear here.")
                trace_file = gr.File(label="⬇️ Agent trace (JSON) — Open Trace", interactive=False)

        run_btn.click(analyze_gpu, inputs=[image_in, profile_in], outputs=[verdict, trace_file])
    return demo


if __name__ == "__main__":
    build_ui().launch()