# =============================================================================
# 🥗 NutriVision - app.py
# Vision Models: nateraw/food | prithivMLmods/Indian-Western-Food-34 | Custom 80-class
# Text AI:       OpenRouter API
# =============================================================================

from flask import Flask, render_template, request, jsonify
from transformers import pipeline, AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import torch
import functools
import os
import re
import requests
import json
from werkzeug.utils import secure_filename

app = Flask(__name__)
app.config["UPLOAD_FOLDER"] = "static/uploads"
app.config["MAX_CONTENT_LENGTH"] = 16 * 1024 * 1024
app.config["ALLOWED_EXTENSIONS"] = {'png', 'jpg', 'jpeg', 'webp'}
os.makedirs(app.config["UPLOAD_FOLDER"], exist_ok=True)

# ============================================================
# 🔑  OPENROUTER CONFIG
# ============================================================
OPENROUTER_API_KEY = "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"
OPENROUTER_URL     = "https://openrouter.ai/api/v1/chat/completions"

CANDIDATE_MODELS = [
    "openai/gpt-4o-mini",
    "mistralai/mistral-7b-instruct:free",
    "google/gemma-2-9b-it:free",
]

# ================================
# 🔹 UTILITIES
# ================================
def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in app.config["ALLOWED_EXTENSIONS"]

def calculate_bmi(height, weight):
    h = height / 100
    return round(weight / (h ** 2), 1)

def get_bmi_category(bmi):
    if bmi < 18.5:   return "Underweight"
    elif bmi < 25.0: return "Normal weight"
    elif bmi < 30.0: return "Overweight"
    else:            return "Obese"

def call_openrouter(prompt, max_tokens=1000):
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type":  "application/json",
        "HTTP-Referer":  "https://nutrivision.ai",
        "X-Title":       "NutriVision",
    }
    for model in CANDIDATE_MODELS:
        print(f"   🔷 Trying model: {model}")
        try:
            payload = {
                "model":       model,
                "messages":    [{"role": "user", "content": prompt}],
                "max_tokens":  max_tokens,
                "temperature": 0.4,
            }
            resp = requests.post(OPENROUTER_URL, headers=headers, json=payload, timeout=45)
            print(f"      HTTP {resp.status_code}")
            if resp.status_code != 200:
                print(f"      ❌ Error: {resp.text[:300]}")
                continue
            content = resp.json().get("choices", [{}])[0].get("message", {}).get("content", "").strip()
            if not content:
                print(f"      ❌ Empty content from {model}")
                continue
            print(f"      ✅ Got {len(content)} chars from {model}")
            return content, model
        except requests.exceptions.Timeout:
            print(f"      ❌ Timeout on {model}")
        except Exception as e:
            print(f"      ❌ Exception on {model}: {e}")
    print("   ❌ All models failed")
    return None, None

# ================================
# 🔹 MODEL 1: nateraw/food
# ================================
@functools.lru_cache(maxsize=1)
def load_food101_classifier():
    print("🔄 [Model 1] Loading nateraw/food …")
    return pipeline("image-classification", model="nateraw/food",
                    device=0 if torch.cuda.is_available() else -1)

# ================================
# 🔹 MODEL 2: Indian-Western-Food-34
# ================================
@functools.lru_cache(maxsize=1)
def load_indian_western_classifier():
    print("🔄 [Model 2] Loading prithivMLmods/Indian-Western-Food-34 …")
    return pipeline("image-classification",
                    model="prithivMLmods/Indian-Western-Food-34",
                    device=0 if torch.cuda.is_available() else -1)

# ================================
# 🔹 MODEL 3: Custom Fine-Tuned
# ================================
@functools.lru_cache(maxsize=1)
def load_custom_model():
    MODEL_PATH = "final_model"
    print("🔄 [Model 3] Loading custom fine-tuned model …")
    try:
        proc = AutoImageProcessor.from_pretrained(MODEL_PATH)
        mdl  = AutoModelForImageClassification.from_pretrained(
            MODEL_PATH,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
        )
        mdl.eval()
        if torch.cuda.is_available():
            mdl = mdl.cuda()
        print("✅ [Model 3] Loaded!")
        return proc, mdl
    except Exception as e:
        print(f"⚠️  [Model 3] Failed: {e}")
        return None, None

# ================================
# 🔹 3-MODEL ENSEMBLE
# ================================
def detect_food(image_path):
    image      = Image.open(image_path).convert('RGB')
    candidates = []

    try:
        preds = load_food101_classifier()(image, top_k=3)
        b = preds[0]
        candidates.append({"food": b['label'].replace('_',' ').title(),
                            "confidence": b['score'], "source": "Food-101"})
        print(f"  ▸ Model 1  {b['label']}  {b['score']*100:.1f}%")
    except Exception as e:
        print(f"  ▸ Model 1 error: {e}")

    try:
        preds = load_indian_western_classifier()(image, top_k=3)
        b = preds[0]
        candidates.append({"food": b['label'].replace('_',' ').title(),
                            "confidence": b['score'], "source": "Indian-Western-34"})
        print(f"  ▸ Model 2  {b['label']}  {b['score']*100:.1f}%")
    except Exception as e:
        print(f"  ▸ Model 2 error: {e}")

    try:
        proc, mdl = load_custom_model()
        if proc and mdl:
            inputs = proc(images=image, return_tensors="pt")
            if torch.cuda.is_available():
                inputs = {k: v.cuda() for k, v in inputs.items()}
            with torch.no_grad():
                logits = mdl(**inputs).logits
            pid  = logits.argmax(-1).item()
            conf = torch.softmax(logits, dim=-1)[0][pid].item()
            lbl  = mdl.config.id2label[pid]
            candidates.append({"food": lbl.replace('_',' ').title(),
                                "confidence": conf, "source": "Custom-80"})
            print(f"  ▸ Model 3  {lbl}  {conf*100:.1f}%")
    except Exception as e:
        print(f"  ▸ Model 3 error: {e}")

    if not candidates:
        return "Unknown Food", 0.0, "No model available"

    winner = max(candidates, key=lambda x: x["confidence"])
    print(f"✅ Winner → {winner['food']}  {winner['confidence']*100:.1f}%  [{winner['source']}]")
    return winner["food"], winner["confidence"], winner["source"]

# ================================
# 🔹 LLM: FULL NUTRITION REPORT
# ================================
def generate_full_report(food_name, age, gender, height, weight,
                          bmi, bmi_category, condition, diet_pref):
    cond_str = condition if condition and condition.lower() != "none" else "None"
    print(f"\n🔶 generate_full_report() → food={food_name}, condition={cond_str}, bmi={bmi_category}")

    prompt = f"""You are a certified nutritionist AI. Return ONLY a raw JSON object — no markdown, no code fences, no explanation, no extra text whatsoever. Start your response with {{ and end with }}.

You are analyzing: {food_name}

User details:
- Age: {age}, Gender: {gender}
- Height: {height}cm, Weight: {weight}kg
- BMI: {bmi} which is {bmi_category}
- Diet: {diet_pref}
- Health condition: {cond_str}

Fill this JSON with REAL, SPECIFIC data for {food_name}. Every field must be specific to {food_name} — never give generic values.

{{
  "nutrition": {{
    "serving_size": "<typical serving size of {food_name}>",
    "calories": "<real calories of {food_name} per serving>",
    "protein": "<real protein of {food_name}>",
    "carbohydrates": "<real carbs of {food_name}>",
    "fat": "<real fat of {food_name}>",
    "fiber": "<real fiber of {food_name}>",
    "sugar": "<real sugar of {food_name}>",
    "sodium": "<real sodium of {food_name}>"
  }},
  "health_benefits": [
    "<benefit 1 specific to {food_name}>",
    "<benefit 2 specific to {food_name}>",
    "<benefit 3 specific to {food_name}>"
  ],
  "portion_advice": "<how much {food_name} should a {age}-year-old {gender} with {bmi_category} BMI and {cond_str} eat>",
  "health_context": "<specific explanation of how {food_name} affects {cond_str} — mention key nutrients and why they matter for {cond_str}>",
  "alternatives": [
    {{"name": "<healthier alternative to {food_name}>", "reason": "<why better for {cond_str} and {bmi_category}>"}},
    {{"name": "<healthier alternative to {food_name}>", "reason": "<why better for {cond_str} and {bmi_category}>"}},
    {{"name": "<healthier alternative to {food_name}>", "reason": "<why better for {cond_str} and {bmi_category}>"}}
  ]
}}"""

    raw, model_used = call_openrouter(prompt, max_tokens=1000)
    if not raw:
        print("⚠️  All LLM calls failed → using fallback")
        return None

    print(f"   Model used: {model_used}")
    print(f"   Raw (first 400 chars): {raw[:400]}")

    try:
        clean = raw.strip()
        clean = re.sub(r"^```[a-zA-Z]*\n?", "", clean)
        clean = re.sub(r"\n?```$", "", clean.strip())
        m = re.search(r'\{.*\}', clean, re.DOTALL)
        if m:
            clean = m.group(0)
        parsed = json.loads(clean)
        print(f"✅ JSON parsed OK — calories={parsed.get('nutrition',{}).get('calories','?')}")
        return parsed
    except Exception as e:
        print(f"⚠️  JSON parse error: {e}")
        print(f"   Raw response: {raw[:600]}")
        return None

# ================================
# 🔹 SHOPPING + DELIVERY URLS
# ================================
def get_shopping_urls(food_item):
    """
    Returns search links for grocery delivery + food delivery platforms.
    Uses each platform's native search URL format.
    """
    raw       = food_item.strip()
    q_pct     = raw.lower().replace(' ', '%20')   # URL percent-encoded
    q_plus    = raw.lower().replace(' ', '+')      # + encoded (Google style)
    q_dash    = raw.lower().replace(' ', '-')      # dash-separated (Swiggy)
    q_zomato  = raw.lower().replace(' ', '%20')    # Zomato uses %20

    return [
        # ── Grocery / delivery platforms ──────────────────────────────
        {
            "platform": "BigBasket",
            "url":      f"https://www.bigbasket.com/ps/?q={q_pct}",
            "emoji":    "🛒",
            "category": "grocery"
        },
        {
            "platform": "Blinkit",
            "url":      f"https://blinkit.com/s/?q={q_pct}",
            "emoji":    "⚡",
            "category": "grocery"
        },
        {
            "platform": "Amazon",
            "url":      f"https://www.amazon.in/s?k={q_plus}+food",
            "emoji":    "📦",
            "category": "grocery"
        },
        {
            "platform": "Flipkart",
            "url":      f"https://www.flipkart.com/search?q={q_pct}",
            "emoji":    "🛍️",
            "category": "grocery"
        },
        # ── Food delivery platforms ────────────────────────────────────
        {
            "platform": "Swiggy",
            "url":      f"https://www.swiggy.com/search?query={q_pct}",
            "emoji":    "🍊",
            "category": "delivery"
        },
        {
            "platform": "Zomato",
            "url":      f"https://www.zomato.com/search?q={q_zomato}",
            "emoji":    "🔴",
            "category": "delivery"
        },
    ]

# ================================
# 🔹 FALLBACK REPORT
# ================================
def fallback_report(food_name="this food"):
    return {
        "nutrition": {
            "serving_size": "1 standard serving (~150g)",
            "calories": "~250 kcal", "protein": "~8g",
            "carbohydrates": "~35g", "fat": "~10g",
            "fiber": "~3g", "sugar": "~5g", "sodium": "~200mg"
        },
        "health_benefits": [
            f"{food_name} provides essential macronutrients for daily energy.",
            "Contains dietary fiber supporting digestive health.",
            "Source of micronutrients important for body functions."
        ],
        "portion_advice": f"Consume 1 standard serving of {food_name} as part of a balanced diet.",
        "health_context": f"Consult a nutritionist for personalised advice about {food_name} and your health goals.",
        "alternatives": [
            {"name": "Steamed Vegetables", "reason": "Low calories, high fiber and nutrients"},
            {"name": "Grilled Chicken",    "reason": "Lean protein, low in saturated fat"},
            {"name": "Fresh Fruit Bowl",   "reason": "Natural sugars with vitamins and antioxidants"}
        ]
    }

# ================================
# 🔹 ROUTES
# ================================
@app.route("/")
def home():
    return render_template("home.html")

@app.route("/analyzer")
def analyzer():
    return render_template("index.html")

@app.route("/about")
def about():
    return render_template("about.html")

@app.route("/analyze", methods=["POST"])
def analyze():
    try:
        if 'image' not in request.files:
            return jsonify({"error": "No image uploaded"}), 400
        image_file = request.files['image']
        if not image_file.filename or not allowed_file(image_file.filename):
            return jsonify({"error": "Invalid file type. Use PNG, JPG, JPEG or WebP."}), 400

        age       = request.form.get("age", "25")
        gender    = request.form.get("gender", "Male")
        height    = float(request.form.get("height", "170"))
        weight    = float(request.form.get("weight", "70"))
        diet_pref = request.form.get("preference", "Vegetarian")
        condition = request.form.get("condition", "None")

        bmi          = calculate_bmi(height, weight)
        bmi_category = get_bmi_category(bmi)

        filename = secure_filename(image_file.filename)
        img_path = os.path.join(app.config["UPLOAD_FOLDER"], filename)
        image_file.save(img_path)

        print("\n" + "="*55)
        print(f"📥 REQUEST: {age}y {gender}, h={height} w={weight}, BMI={bmi} ({bmi_category})")
        print(f"   condition={condition}, diet={diet_pref}")

        print("\n━━━ 3-MODEL ENSEMBLE ━━━")
        food_name, confidence, detection_source = detect_food(img_path)

        print("\n━━━ LLM NUTRITION REPORT ━━━")
        report = generate_full_report(
            food_name, age, gender, height, weight,
            bmi, bmi_category, condition, diet_pref
        )

        if report is None:
            print("⚠️  Using FALLBACK")
            report = fallback_report(food_name)

        alternatives = [
            {"name": a["name"], "reason": a["reason"],
             "urls": get_shopping_urls(a["name"])}
            for a in report.get("alternatives", [])
        ]

        return jsonify({
            "food":             food_name,
            "confidence":       f"{confidence * 100:.1f}%",
            "detection_source": detection_source,
            "bmi":              bmi,
            "bmi_category":     bmi_category,
            "nutrition":        report.get("nutrition", {}),
            "health_benefits":  report.get("health_benefits", []),
            "portion_advice":   report.get("portion_advice", "1 standard serving"),
            "health_context":   report.get("health_context", ""),
            "alternatives":     alternatives,
        })

    except Exception as e:
        import traceback; traceback.print_exc()
        return jsonify({"error": f"Analysis failed: {str(e)}"}), 500

if __name__ == "__main__":
    print("🚀 NutriVision starting…")
    print(f"🎮 GPU: {torch.cuda.is_available()}")
    print(f"🔑 OpenRouter key: {OPENROUTER_API_KEY[:18]}...")
    print(f"🤖 Model priority: {CANDIDATE_MODELS}")
    app.run(host="0.0.0.0", port=7860)