import gradio as gr
import os
import json
import requests
from PIL import Image
from difflib import get_close_matches
from functools import lru_cache
from collections import Counter

# =====================
# ENV
# =====================
ROBOFLOW_API_KEY = os.getenv("ROBOFLOW_API_KEY")

# =====================
# LOAD DB
# =====================
with open("nutrition_db.json", "r") as f:
    NUTRITION_DB = json.load(f)

print("✅ Loaded DB:", len(NUTRITION_DB))

# =====================
# COUNTABLE FOODS (NEW)
# =====================
COUNTABLE_FOODS = {
    "roti": {"calories": 120, "protein": 3, "carbs": 20, "fat": 2},
    "bread": {"calories": 80, "protein": 3, "carbs": 15, "fat": 1},
    "samosa": {"calories": 260, "protein": 5, "carbs": 30, "fat": 14},
    "gulab jamun": {"calories": 150, "protein": 2, "carbs": 30, "fat": 5},
    "laddu": {"calories": 180, "protein": 3, "carbs": 25, "fat": 8},
    "idli": {"calories": 60, "protein": 2, "carbs": 12, "fat": 1},
    "vada": {"calories": 150, "protein": 4, "carbs": 20, "fat": 8}
}

# =====================
# NORMALIZATION (FIXED)
# =====================
def normalize_food_name(name):
    name = name.lower()

    mapping = {
        "chapati": "roti",
        "omellete": "omelette",
        "fried rice": "rice",
        "plain rice": "rice"
    }

    return mapping.get(name, name)

# =====================
# FIND MATCH
# =====================
def find_food(name):
    if name in NUTRITION_DB:
        return name

    matches = get_close_matches(name, NUTRITION_DB.keys(), n=1, cutoff=0.6)
    return matches[0] if matches else None

# =====================
# FALLBACK
# =====================
def estimate_unknown_food(grams):
    return {
        "calories": round(1.5 * grams, 2),
        "protein": round(0.05 * grams, 2),
        "carbs": round(0.2 * grams, 2),
        "fat": round(0.05 * grams, 2),
    }

# =====================
# CACHE
# =====================
@lru_cache(maxsize=1000)
def compute_nutrition_cached(food_key, grams):
    base = NUTRITION_DB[food_key]
    factor = grams / 100

    return {
        "calories": round(base["calories"] * factor, 2),
        "protein": round(base["protein"] * factor, 2),
        "carbs": round(base["carbs"] * factor, 2),
        "fat": round(base["fat"] * factor, 2),
    }

# =====================
# GET NUTRITION
# =====================
def get_nutrition(dish, grams):
    dish = normalize_food_name(dish)
    food_key = find_food(dish)

    if not food_key:
        return estimate_unknown_food(grams)

    return compute_nutrition_cached(food_key, grams)

# =====================
# QUANTITY ESTIMATION
# =====================
def estimate_quantity(pred):
    width = pred.get("width", 0)
    height = pred.get("height", 0)

    area = width * height
    ratio = area / (640 * 640)

    grams = 150 + (ratio * 300)
    return round(grams, 1)

# =====================
# ROBOFLOW DETECTION
# =====================
def detect(image_path):
    api_key = os.getenv("ROBOFLOW_API_KEY")

    if not api_key:
        return "❌ API KEY NOT FOUND"

    url = "https://detect.roboflow.com/almost-final/2"

    try:
        with open(image_path, "rb") as f:
            response = requests.post(
                url,
                files={"file": f},
                params={"api_key": api_key},
                timeout=15
            )

        if response.status_code != 200:
            return f"❌ Roboflow Error: {response.text}"

        data = response.json()
        return data.get("predictions", [])

    except Exception as e:
        return f"❌ Request Failed: {str(e)}"

# =====================
# MAIN FUNCTION (HYBRID)
# =====================
def analyze_image(image):
    if image is None:
        return "Please upload an image"

    path = "temp.jpg"
    image.save(path)

    preds = detect(path)

    if isinstance(preds, str):
        return preds

    if len(preds) == 0:
        return "❌ No food detected"

    # GROUP ITEMS
    dish_counts = Counter([
        normalize_food_name(p.get("class", "unknown")) for p in preds
    ])

    output = ""
    total = {"calories": 0, "protein": 0, "carbs": 0, "fat": 0}

    for dish, count in dish_counts.items():

        if dish == "unknown":
            continue

        # =========================
        # COUNTABLE FOODS
        # =========================
        if dish in COUNTABLE_FOODS:
            base = COUNTABLE_FOODS[dish]

            nutrition = {
                "calories": base["calories"] * count,
                "protein": base["protein"] * count,
                "carbs": base["carbs"] * count,
                "fat": base["fat"] * count,
            }

            output += f"🍽️ {dish} ({count} pcs)\n"

        # =========================
        # WEIGHT-BASED FOODS
        # =========================
        else:
            relevant_preds = [
                p for p in preds if normalize_food_name(p.get("class")) == dish
            ]

            grams_list = [estimate_quantity(p) for p in relevant_preds]
            grams = sum(grams_list)

            nutrition = get_nutrition(dish, grams)

            output += f"🍽️ {dish}\n"
            output += f"📏 {grams:.1f} g\n"

        # =========================
        # OUTPUT
        # =========================
        output += f"🔥 {nutrition['calories']} kcal\n"
        output += f"💪 Protein: {nutrition['protein']} g\n"
        output += f"🍞 Carbs: {nutrition['carbs']} g\n"
        output += f"🧈 Fat: {nutrition['fat']} g\n"
        output += "-"*30 + "\n"

        for k in total:
            total[k] += nutrition[k]

    # TOTAL
    output += "\n🧾 TOTAL:\n"
    output += f"🔥 Calories: {round(total['calories'],2)}\n"
    output += f"💪 Protein: {round(total['protein'],2)} g\n"
    output += f"🍞 Carbs: {round(total['carbs'],2)} g\n"
    output += f"🧈 Fat: {round(total['fat'],2)} g\n"

    return output

# =====================
# UI
# =====================
demo = gr.Interface(
    fn=analyze_image,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="🍽️ AI Nutritionist",
    description="Upload food image to get calories & macros"
)

# =====================
# RUN
# =====================
if __name__ == "__main__":
    demo.launch(ssr_mode=False)