import base64
import io
import json
import os

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

API_KEY = os.environ.get("KHANA_API_KEY", "")
MODEL_ID = "Qwen/Qwen3-VL-4B-Instruct"

# Load to CPU at startup; moved to CUDA inside @spaces.GPU per-call
model = Qwen3VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
)
processor = AutoProcessor.from_pretrained(MODEL_ID)

FOOD_PROMPT = """You are a nutrition expert. Analyze this food image and return ONLY raw JSON, no markdown fences.

For each food item visible return:
{{
  "is_real_food": true,
  "is_harmful": false,
  "harmful_reason": "",
  "items": [
    {{
      "name": "item name",
      "quantity": "e.g. 1 cup / 200g / 1 piece",
      "calories": 250,
      "protein_g": 10.0,
      "carbs_g": 30.0,
      "fat_g": 8.0,
      "saturated_fat_g": 3.0,
      "sugars_g": 5.0,
      "fiber_g": 2.0,
      "sodium_mg": 400
    }}
  ],
  "total_calories": 250,
  "notes": "one short line: label read / estimated from similar product / assumptions"
}}

RULES:
- For packaged food, read the Nutrition Facts label if visible; otherwise use known product values.
- If the product is not recognisable, speculate from the closest known product and say so in notes.
- Estimate for the full visible serving unless the user specifies otherwise.
- Judge alcohol from packaging, label, or context -- NOT colour. Black coffee, cranberry juice, pomegranate juice, kokum, rooh afza are NOT alcohol.
- If the image shows alcohol, tobacco, a vape, or drugs: return {{"is_real_food": true, "is_harmful": true, "harmful_reason": "what it is", "items": [], "total_calories": 0, "notes": "harmful"}}
- If the image is NOT food at all (screenshot, ad, cartoon): return {{"is_real_food": false, "items": [], "total_calories": 0, "notes": "reason"}}
"""


def _extract_json(text: str) -> str:
    text = text.strip()
    if "```" in text:
        for part in text.split("```"):
            part = part.strip()
            if part.startswith("json"):
                part = part[4:].strip()
            if part.startswith("{"):
                return part
    return text


@spaces.GPU(duration=120)
def analyze(image_b64: str, extra_notes: str, api_key: str) -> str:
    if API_KEY and api_key.strip() != API_KEY:
        return json.dumps({"error": "unauthorized"})

    try:
        img_bytes = base64.b64decode(image_b64)
        image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
    except Exception as e:
        return json.dumps({"error": f"image decode failed: {e}"})

    RAW_PREFIX = "RAW_PROMPT:"
    raw_mode = extra_notes.startswith(RAW_PREFIX)
    if raw_mode:
        # benchmark mode: use extra_notes verbatim as the whole prompt, no JSON wrapping
        prompt = extra_notes[len(RAW_PREFIX):]
    else:
        prompt = FOOD_PROMPT
        if extra_notes:
            prompt += f"\nUser notes: {extra_notes}"

    messages = [{
        "role": "user",
        "content": [
            {"type": "image", "image": image},
            {"type": "text", "text": prompt},
        ],
    }]

    try:
        inputs = processor.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_dict=True,
            return_tensors="pt",
        ).to("cuda")

        model.to("cuda")
        with torch.no_grad():
            output_ids = model.generate(**inputs, max_new_tokens=(30 if raw_mode else 900), do_sample=False)

        generated = output_ids[0][inputs.input_ids.shape[1]:]
        raw = processor.decode(generated, skip_special_tokens=True)
        return raw if raw_mode else _extract_json(raw)
    except Exception as e:
        import traceback
        return json.dumps({"error": f"{type(e).__name__}: {e}", "trace": traceback.format_exc()[-1500:]})


demo = gr.Interface(
    fn=analyze,
    inputs=[
        gr.Textbox(label="image_b64", lines=1, placeholder="base64-encoded image"),
        gr.Textbox(label="extra_notes", lines=1, placeholder="e.g. 2 rotis, no oil"),
        gr.Textbox(label="api_key", lines=1, placeholder="KHANA_API_KEY value"),
    ],
    outputs=gr.Textbox(label="result_json"),
    title="KhanaVision",
    description="Qwen3-VL-4B-Instruct food analysis endpoint for Khana Tracker",
    allow_flagging="never",
)

demo.launch()