app.py

"""
NutriLens - Food Health Impact Analyzer
Gradio Build Small Hackathon (June 2026)
"""

import os
import json
import re
import base64
import io
import time
import gradio as gr
from PIL import Image
from dotenv import load_dotenv
from huggingface_hub import InferenceClient

load_dotenv()

from src.nutrition import lookup_ingredients
from src.literature import lookup_literature, format_citation
from src.prompts import IDENTIFY_PROMPT, build_analysis_prompt, HEALTH_GOALS, AUDIENCES

# ---- Configuration ----
MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen3.6-27B")
API_BASE = os.environ.get("API_BASE", None)
HF_TOKEN = os.environ.get("HF_TOKEN", None)

client = InferenceClient(
    model=API_BASE or MODEL_ID,
    token=HF_TOKEN,
    timeout=300,
)

# ---- Custom CSS ----
CUSTOM_CSS = """
.nutrilens-report h2 {
    color: #2d8659;
    border-bottom: 2px solid #2d8659;
    padding-bottom: 6px;
    margin-top: 24px;
}
.nutrilens-report h3 {
    color: #5b6abf;
    margin-top: 20px;
}
.summary-card {
    background: linear-gradient(135deg, #e8f5e9 0%, #e3f2fd 100%);
    border-left: 4px solid #2d8659;
    border-radius: 8px;
    padding: 16px 20px;
    margin: 12px 0;
    color: #1a3a2a;
}
.dark .summary-card {
    background: linear-gradient(135deg, #1b3a2a 0%, #1a2a3a 100%);
    color: #c8e6c9;
}
.tip-card {
    background: #fff8e1;
    border-left: 4px solid #f9a825;
    border-radius: 8px;
    padding: 16px 20px;
    margin: 12px 0;
    color: #4a3800;
}
.dark .tip-card {
    background: #2a2510;
    color: #ffe082;
}
.watch-out-label {
    color: #c62828;
}
.dark .watch-out-label {
    color: #ff6b6b;
}
.disclaimer-box {
    background: #fce4ec;
    border-left: 4px solid #e53935;
    border-radius: 8px;
    padding: 12px 16px;
    margin: 16px 0;
    color: #4a0e0e;
    font-size: 0.9em;
}
.dark .disclaimer-box {
    background: #2a1010;
    color: #ef9a9a;
}
"""


def call_model(messages: list, max_tokens: int = 1024, retries: int = 2,
               extract_answer: bool = True, markers: tuple = None) -> str:
    """Call model with timeout handling and retry.
    Set extract_answer=False for ingredient ID (has its own parser).
    `markers`, if given, is a (start, end) pair of sentinel lines the
    prompt asked the model to wrap its final answer in - tried before
    any heuristic extraction since it's deterministic."""
    for attempt in range(retries + 1):
        try:
            response = client.chat_completion(
                model=MODEL_ID if not API_BASE else None,
                messages=messages,
                max_tokens=max_tokens,
                temperature=0.3,
            )
            msg = response.choices[0].message
            content = msg.content
            reasoning = getattr(msg, "reasoning", None) or ""

            # The model may put its final answer in `content`, in
            # `reasoning`, or wrap it with the sentinel markers in either -
            # whichever field actually has text is the one to look at first.
            text = content if content is not None else reasoning
            if not text:
                return "The model returned an empty response. Please try again."

            between = _extract_between_markers(text, *markers) if markers else None
            if between is not None:
                content = between
            elif content is None:
                content = _extract_answer_from_reasoning(reasoning) if extract_answer else reasoning
            # else: keep msg.content as-is - it's real content with no markers needed

            content = re.sub(r"<think>.*?</think>", "", content, flags=re.DOTALL).strip()
            return content if content else "Model returned empty content."

        except Exception as e:
            error_str = str(e)
            if ("504" in error_str or "timeout" in error_str.lower()) and attempt < retries:
                wait = 3 * (attempt + 1)
                print(f"Timeout (attempt {attempt+1}/{retries+1}), retrying in {wait}s...")
                time.sleep(wait)
                continue
            print(f"Model call error: {e}")
            if "504" in error_str:
                return ("The model server timed out. This usually happens with long "
                        "ingredient lists. Try with fewer ingredients (5-8 at a time).")
            return f"Error calling model: {e}"
    return "All retries failed. Please try again later."


def _extract_between_markers(text: str, start: str, end: str) -> str | None:
    """Return the text between two sentinel marker lines, or None if no
    substantial match is found. The prompt asks the model to wrap its
    final answer in these markers - but while thinking, the model often
    also *mentions* the marker format (e.g. "wrap the answer between
    @@@REPORT_START@@@ and @@@REPORT_END@@@"), which produces a tiny,
    bogus match. The real final answer is always much longer than any
    incidental mention, so among all matches we take the longest one
    that clears a minimum length."""
    pattern = re.escape(start) + r"\s*(.*?)\s*" + re.escape(end)
    matches = [m.strip() for m in re.findall(pattern, text, re.DOTALL)]
    matches = [m for m in matches if len(m) > 60]
    if matches:
        return max(matches, key=len)

    # The model can also get cut off mid-answer (hits max_tokens before
    # emitting the end marker). In that case there's no complete pair, but
    # the last start-marker occurrence is still where the real answer
    # begins - take everything after it rather than leaking the raw
    # "@@@REPORT_START@@@" line to the user.
    starts = [m.end() for m in re.finditer(re.escape(start), text)]
    if starts:
        tail = text[starts[-1]:].strip()
        if len(tail) > 200:
            return tail
    return None


def _extract_answer_from_reasoning(reasoning: str) -> str:
    """
    When Qwen3.6 thinks, the reasoning field contains both the
    internal chain-of-thought AND the final formatted answer.
    This function extracts just the answer.
    """
    # Look for markdown headings at the START of a line (not inline mentions).
    # The model's self-checks mention headings inline like:
    #   'Is "## What's on your plate" present? Yes.'
    # But the actual answer has them at line start:
    #   '\n## What's on your plate\n'
    markers = [
        r"\n## What.s on your plate",
        r"\n## What.s on Your Plate",
        r"\n## Overall Meal",
        r"\n## Overall Assessment",
        r"\n## Summary",
    ]
    for pattern in markers:
        matches = list(re.finditer(pattern, reasoning, re.IGNORECASE))
        if matches:
            # Use the LAST match that's at a line start
            idx = matches[-1].start() + 1  # +1 to skip the \n
            answer = reasoning[idx:]
            # Trim trailing thinking artifacts
            for end_marker in ["✅", "[Done]", "[Output Generation]",
                               "Self-Correction", "Output matches"]:
                end_idx = answer.rfind(end_marker)
                if end_idx > 0 and end_idx > len(answer) * 0.6:
                    answer = answer[:end_idx].strip()
            if len(answer) > 50:
                return answer.strip()

    # Fallback: look for the last block of markdown-formatted text
    # by finding consecutive lines starting with ## or ### or - or *
    lines = reasoning.split('\n')
    best_start = None
    for i, line in enumerate(lines):
        if line.strip().startswith('## ') and not '`' in line and '?' not in line:
            # This looks like a real heading, not a self-check
            if best_start is None:
                best_start = i
    if best_start is not None:
        answer = '\n'.join(lines[best_start:])
        if len(answer) > 50:
            return answer.strip()

    # For ingredient identification: try to find JSON array
    # Look for the largest JSON array (not tiny ones like [1])
    json_matches = re.findall(r'\[("[^"]+?"(?:\s*,\s*"[^"]+?")*)\]', reasoning)
    if json_matches:
        longest = max(json_matches, key=len)
        return f'[{longest}]'

    # Last resort: return the last 30% of reasoning
    cutoff = int(len(reasoning) * 0.7)
    return reasoning[cutoff:].strip()


def image_to_data_url(img: Image.Image) -> str:
    buf = io.BytesIO()
    if img.mode == "RGBA":
        img = img.convert("RGB")
    max_dim = 1024
    if max(img.size) > max_dim:
        img.thumbnail((max_dim, max_dim), Image.LANCZOS)
    img.save(buf, format="JPEG", quality=80)
    b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
    return f"data:image/jpeg;base64,{b64}"


def extract_ingredients_from_text(text: str) -> list[str]:
    NOISE = {
        "zutaten", "ingredients", "ingredienten", "ingrédients", "sastojci",
        "contains", "kann auch", "may contain", "enthält", "contient",
        "allergens", "allergenen", "nutrition", "nährwerte",
        "analyze", "image", "ingredient", "label", "user", "step", "json",
        "the", "and", "oder", "und", "et", "i", "a", "an",
    }

    # Phrases that only show up when the model is restating its own
    # instructions (while thinking) rather than naming a food - reject
    # any item that contains one of these, regardless of language/case.
    INSTRUCTION_PHRASES = (
        "actual food", "list only", "json array", "final answer",
        "marker", "format", "translate", "do not include", "do not repeat",
    )

    def is_food(item: str) -> bool:
        item_lower = item.lower().strip()
        if len(item_lower) < 2 or len(item_lower) > 80:
            return False
        if item_lower in NOISE:
            return False
        if any(item_lower.startswith(n) for n in ["zutaten", "may contain", "kann auch"]):
            return False
        if any(p in item_lower for p in INSTRUCTION_PHRASES):
            return False
        return True

    # Try every bracketed array in the text (the model may mention an
    # example array while thinking before producing the real, final one)
    # and keep whichever yields the most valid food items - the genuine
    # final list is reliably the longest, most complete one.
    best = []
    for candidate in re.findall(r'\[.*?\]', text, re.DOTALL):
        try:
            items = json.loads(candidate)
        except json.JSONDecodeError:
            continue
        if not isinstance(items, list):
            continue
        foods = [str(i).strip() for i in items if is_food(str(i))]
        if len(foods) > len(best):
            best = foods
    if best:
        return list(dict.fromkeys(best))

    quoted = re.findall(r'"([^"]+)"', text)
    if len(quoted) >= 2:
        foods = [q for q in quoted if is_food(q)]
        if foods:
            return list(dict.fromkeys(foods))

    arrow_matches = re.findall(r'->\s*([a-zA-Z][a-zA-Z\s,]+?)(?:\n|$)', text)
    if arrow_matches:
        foods = [m.strip().rstrip(',').strip() for m in arrow_matches if is_food(m.strip())]
        if foods:
            return list(dict.fromkeys(foods))

    parts = [p.strip() for p in text.split(',') if p.strip()]
    foods = [p for p in parts if is_food(p)]
    if foods:
        return list(dict.fromkeys(foods))[:20]

    return [text[:100]]


def format_report(raw_report: str, nutrition_fails: int, lit_fails: int,
                  total_ingredients: int) -> str:
    """Post-process the model's markdown into styled HTML."""
    report = raw_report

    # The model doesn't always put a line break before "**Watch out:**" -
    # it sometimes lands mid-sentence on the same line as the last "Good
    # stuff" bullet, which makes Markdown swallow it into that list item.
    # Force it onto its own paragraph and color it so it stands out.
    report = re.sub(
        r"\s*\*\*Watch out:\*\*",
        '\n\n<strong class="watch-out-label">Watch out:</strong>',
        report,
    )

    # Cut the tips section out first (wherever the model placed it) so we
    # can re-insert it right before the ingredient breakdown instead of at
    # the end - the user wants tips to appear up front, near the summary.
    tips_match = re.search(
        r"##\s*Tips?\s*\n(.*?)(?=\n##|\Z)",
        report, re.DOTALL | re.IGNORECASE
    )
    tips_html = ""
    if tips_match:
        tips_text = tips_match.group(1).strip()
        tips_html = (f"## Tips\n\n"
                     f'<div class="tip-card">\n\n{tips_text}\n\n</div>\n\n')
        report = report[:tips_match.start()] + report[tips_match.end():]

    # Wrap the summary section in a card
    summary_match = re.search(
        r"(##\s*What.s on your plate\s*\n)(.*?)(?=\n##|\n###|\Z)",
        report, re.DOTALL | re.IGNORECASE
    )
    if summary_match:
        summary_text = summary_match.group(2).strip()
        styled = (f"## What's on your plate\n\n"
                  f'<div class="summary-card">\n\n{summary_text}\n\n</div>\n\n')
        report = report[:summary_match.start()] + styled + report[summary_match.end():]

    # Re-insert tips right before the first ingredient heading (### ...),
    # which immediately follows the summary card.
    if tips_html:
        first_heading = re.search(r"\n###\s", report)
        if first_heading:
            insert_at = first_heading.start() + 1
            report = report[:insert_at] + tips_html + "\n" + report[insert_at:]
        else:
            report += "\n\n" + tips_html

    # Add disclaimer card
    disclaimer = (
        '<div class="disclaimer-box">'
        '⚠️ <strong>This is not medical advice.</strong> '
        'Always talk to a doctor or nutritionist before changing your diet, '
        'especially if you have health conditions, allergies, or take medication.'
        '</div>'
    )

    # Source info
    source = "\n\n---\n*Data: "
    if nutrition_fails == 0:
        source += "USDA FoodData Central"
    elif nutrition_fails < total_ingredients:
        source += "USDA (partial)"
    else:
        source += "Model knowledge"
    source += " + PubMed"
    if lit_fails > 0:
        source += " (partial)"
    source += f" | Model: {MODEL_ID}*"

    return report + "\n\n" + disclaimer + source


def identify_ingredients(image, text_input):
    if image is not None:
        gr.Info("Reading image with AI... this can take 15-30 seconds.")
        data_url = image_to_data_url(image)
        messages = [{
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {"url": data_url}},
                {"type": "text", "text": IDENTIFY_PROMPT},
            ],
        }]
        # Generous budget: with thinking mode on, the model works through
        # the label (translating, deduplicating) before writing the final
        # marker-wrapped array - too small a cap truncates it mid-thought,
        # leaving only messy draft arrays (mixed German/English) behind.
        raw = call_model(
            messages, max_tokens=3000, extract_answer=False,
            markers=("@@@INGREDIENTS_START@@@", "@@@INGREDIENTS_END@@@"),
        )
        ingredients = extract_ingredients_from_text(raw)
        gr.Info(f"Found {len(ingredients)} ingredients. Review and edit if needed.")
        return ", ".join(ingredients)

    elif text_input and text_input.strip():
        items = [i.strip() for i in text_input.replace("\n", ",").split(",") if i.strip()]
        return ", ".join(items)

    return ""


def run_analysis(ingredients_text, health_goal, audience, progress=gr.Progress()):
    if not ingredients_text or not ingredients_text.strip():
        return "Please identify ingredients first.", ""

    ingredients = [i.strip() for i in ingredients_text.split(",") if i.strip()]
    if not ingredients:
        return "No ingredients to analyze.", ""

    # Cap tokens based on ingredient count to avoid timeouts.
    # Thinking + a sentinel-wrapped final answer use more tokens than a
    # bare answer, so the budget is a bit larger than before.
    # The model spends a fairly fixed chunk of its budget on thinking
    # regardless of list length, so short lists need a floor too - without
    # it, thinking alone can exhaust the budget and truncate the report.
    max_tok = min(12000, max(6000, 2000 + len(ingredients) * 600))

    try:
        progress(0.05, desc="Looking up nutritional data...")
        nutrition_data, nutrition_fails = lookup_ingredients(ingredients)

        progress(0.35, desc="Searching scientific literature...")
        goal_key = health_goal if health_goal in HEALTH_GOALS else "General"
        lit_data, lit_fails = lookup_literature(
            ingredients, health_goal=goal_key, papers_per=2
        )

        progress(0.6, desc="Generating health report... this can take 30-90 seconds.")
        prompt = build_analysis_prompt(
            nutrition_data, lit_data, health_goal=goal_key,
            audience=audience,
            nutrition_failures=nutrition_fails,
            literature_failures=lit_fails,
        )
        raw_report = call_model(
            [{"role": "user", "content": prompt}],
            max_tokens=max_tok,
            markers=("@@@REPORT_START@@@", "@@@REPORT_END@@@"),
        )

        progress(0.95, desc="Formatting report...")
        report = format_report(raw_report, nutrition_fails, lit_fails, len(ingredients))

        # Citations
        all_citations = []
        for papers in lit_data.values():
            for p in papers:
                c = format_citation(p)
                if c not in all_citations:
                    all_citations.append(c)

        citations = ""
        if all_citations:
            citations = "**References:**\n\n"
            for i, c in enumerate(all_citations, 1):
                citations += f"{i}. {c}\n\n"

        return report, citations

    except Exception as e:
        import traceback
        traceback.print_exc()
        return f"Something went wrong: {e}", ""


def _start_identify_loading():
    return gr.update(value="⏳ Reading...", interactive=False)


def _stop_identify_loading():
    return gr.update(value="1. Identify ingredients", interactive=True)


def _start_analyze_loading():
    placeholder = (
        "_Generating your health report - looking up nutrition data, "
        "searching scientific literature, and writing up the analysis. "
        "This can take 30-90 seconds..._"
    )
    return (
        gr.update(value="⏳ Analyzing... (30-90s)", interactive=False),
        placeholder,
        "",
    )


def _stop_analyze_loading():
    return gr.update(value="2. Analyze health impact", interactive=True)


# ---- Gradio UI ----

with gr.Blocks(
    title="NutriLens",
    theme=gr.themes.Soft(
        primary_hue="green",
        secondary_hue="blue",
    ),
    css=CUSTOM_CSS,
) as demo:
    gr.Markdown("""
# 🔬 NutriLens
**Upload a food photo or type ingredients, then get a clear health
breakdown backed by real data and scientific research.**

Works with food labels in any language.
    """)

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(
                label="Food photo or ingredient label",
                type="pil",
                sources=["upload", "webcam", "clipboard"],
            )
            text_input = gr.Textbox(
                label="Or type ingredients (comma-separated)",
                placeholder="chicken breast, brown rice, broccoli, olive oil",
                lines=2,
            )
            identify_btn = gr.Button(
                "1. Identify ingredients", variant="secondary", size="lg",
            )

        with gr.Column(scale=2):
            ingredients_box = gr.Textbox(
                label="Identified ingredients (review and edit before analyzing)",
                placeholder="Ingredients will appear here. Edit them if needed, then click Analyze.",
                lines=2,
                interactive=True,
            )
            with gr.Row():
                health_goal = gr.Dropdown(
                    label="Health focus",
                    choices=list(HEALTH_GOALS.keys()),
                    value="General",
                    scale=2,
                )
                audience = gr.Radio(
                    label="Explanation level",
                    choices=list(AUDIENCES.keys()),
                    value="Everyone",
                    scale=2,
                )
            analyze_btn = gr.Button(
                "2. Analyze health impact", variant="primary", size="lg",
            )
            report_out = gr.Markdown(
                label="Health report",
                elem_classes=["nutrilens-report"],
            )
            citations_out = gr.Markdown(label="References")

    identify_btn.click(
        fn=_start_identify_loading,
        outputs=[identify_btn],
    ).then(
        fn=identify_ingredients,
        inputs=[image_input, text_input],
        outputs=[ingredients_box],
    ).then(
        fn=_stop_identify_loading,
        outputs=[identify_btn],
    )

    analyze_btn.click(
        fn=_start_analyze_loading,
        outputs=[analyze_btn, report_out, citations_out],
    ).then(
        fn=run_analysis,
        inputs=[ingredients_box, health_goal, audience],
        outputs=[report_out, citations_out],
    ).then(
        fn=_stop_analyze_loading,
        outputs=[analyze_btn],
    )

    gr.Markdown("""
---
⚠️ **NutriLens is not a substitute for professional medical advice.**
Always consult a doctor or registered nutritionist before making dietary
changes, especially if you have health conditions, allergies, or take
medication.

*Data: USDA FoodData Central + PubMed | Model: ≤32B params |
Built for the [Build Small Hackathon](https://huggingface.co/build-small-hackathon)*
    """)

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