import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image
import json
from datetime import datetime
import tempfile
import os
from gtts import gTTS
from fpdf import FPDF
from health_logic import generate_advice
from tensorflow.keras.applications.efficientnet import preprocess_input as efficientnet_preprocess

# Load model
model = tf.keras.models.load_model("food_vision_model.keras")

class_names = [
    "akara", "banga_soup", "egusi_soup", "jollof_rice", "moi_moi",
    "nkwobi", "okpa", "suya", "tuwo", "yam_porridge"
]

# Load food metadata
with open("food_info.json", "r") as f:
    food_info = json.load(f)

# Language map for TTS
LANG_CODE = {
    "English": "en",
    "Hausa": "en",  # gTTS doesn't support Hausa properly
    "Yoruba": "en",
    "Igbo": "en"
}

# Risk level helper
def get_risk_level(score):
    if score <= 30:
        return "🟢 Low Risk"
    elif score <= 70:
        return "🟡 Medium Risk"
    else:
        return "🔴 High Risk"

# PDF class with Unicode support
class UnicodePDF(FPDF):
    def header(self):
        self.set_font("DejaVu", "B", 16)
        self.cell(0, 10, "HoodHealth Pro+ Report", ln=True)

    def footer(self):
        self.set_y(-15)
        self.set_font("DejaVu", "I", 8)
        self.cell(0, 10, f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M')}", align="C")

def generate_pdf(info, advice, risk_score, risk_level, flags, conditions, lang, display_name, image_path):
    pdf = UnicodePDF()
    pdf.add_font("DejaVu", "", "DejaVuSans.ttf", uni=True)
    pdf.add_font("DejaVu", "B", "DejaVuSans.ttf", uni=True)
    pdf.add_font("DejaVu", "I", "DejaVuSans.ttf", uni=True)
    pdf.add_page()
    pdf.set_font("DejaVu", "", 12)

    pdf.image(image_path, x=10, y=40, w=60)
    pdf.set_xy(75, 40)
    pdf.multi_cell(0, 10, f"Food: {display_name}\nEthnicity: {info['ethnicity']}\nIngredients: {info['ingredients']}")

    pdf.ln(35)
    pdf.multi_cell(0, 10, f"Calories: {info['calories']} kcal\nCarbs: {info['carbs']}g\nProtein: {info['protein']}g\nFat: {info['fat']}g")
    pdf.ln(5)
    pdf.multi_cell(0, 10, f"Diet Type: {info['diet_type']}\nSubstitute: {info.get('substitute', 'None')}")
    pdf.ln(10)

    pdf.set_font("DejaVu", "B", 14)
    pdf.cell(0, 10, "Health Analysis", ln=True)
    pdf.set_font("DejaVu", "", 12)
    pdf.multi_cell(0, 10, f"Conditions: {', '.join(conditions)}\nRisk Score: {risk_score}% ({risk_level})\nRisk Factors: {', '.join(flags) if flags else 'None'}\nAdvice: {advice}")
    pdf.ln(5)
    pdf.set_font("DejaVu", "I", 10)
    pdf.multi_cell(0, 10, "*Note: This advice is not a medical diagnosis.*")

    path = os.path.join(tempfile.gettempdir(), "report.pdf")
    pdf.output(path)
    return path

# Main function
def classify_food(image, conditions, language):
    if not conditions:
        return "⚠️ Please select at least one health condition.", None, None, None

    # Preprocess image
    image = image.convert("RGB").resize((224, 224))
    img = tf.keras.preprocessing.image.img_to_array(image)
    img = efficientnet_preprocess(img)
    img = np.expand_dims(img, axis=0)

    # Predict
    preds = model.predict(img)[0]
    idx = np.argmax(preds)
    predicted_class = class_names[idx]
    confidence = round(float(preds[idx]) * 100, 2)

    info = food_info[predicted_class]
    display_name = info.get("display_name", predicted_class.replace("_", " ").title())

    # Advice logic
    advice, risk_score, flags = generate_advice(info, conditions)
    risk_level = get_risk_level(risk_score)

    # Save audio
    audio_path = os.path.join(tempfile.gettempdir(), "tts.mp3")
    tts = gTTS(text=advice, lang=LANG_CODE.get(language, "en"))
    tts.save(audio_path)

    # Save image
    image_path = os.path.join(tempfile.gettempdir(), "upload.jpg")
    image.save(image_path)

    # Generate PDF
    pdf_path = generate_pdf(info, advice, risk_score, risk_level, flags, conditions, language, display_name, image_path)

    # Output summary
    result = (
        f"🍽️ Food: {display_name}\n"
        f"🌍 Ethnicity: {info['ethnicity']}\n"
        f"🥦 Ingredients: {info['ingredients']}\n"
        f"🔥 Calories: {info['calories']} kcal\n"
        f"🍞 Carbs: {info['carbs']}g\n"
        f"🥩 Protein: {info['protein']}g\n"
        f"🧈 Fat: {info['fat']}g\n"
        f"🌱 Diet Type: {info['diet_type']}\n"
        f"🔁 Substitute: {info.get('substitute', 'None')}\n\n"
        f"📈 Confidence: {confidence}%\n"
        f"📊 Risk Score: {risk_score}% ({risk_level})\n"
        f"⚠️ Risk Factors: {', '.join(flags) if flags else 'None'}\n"
        f"✅ Advice: {advice}\n\n"
        f"📝 *Generated by HoodHealth Pro+.*"
    )

    return result, audio_path, f"🗣️ Language: {language}", pdf_path

# Gradio interface
interface = gr.Interface(
    fn=classify_food,
    inputs=[
        gr.Image(type="pil", label="Upload Food Image"),
        gr.CheckboxGroup(label="Select Health Conditions", choices=[
            "Normal", "Diabetic", "Hypertensive", "Weight Loss", "Malnourished", "Pregnant/Nursing", "Cholesterol Watch"]),
        gr.Dropdown(label="Language for TTS", choices=["English", "Hausa", "Yoruba", "Igbo"], value="English")
    ],
    outputs=[
        gr.Textbox(label="Prediction and Advice"),
        gr.Audio(label="Hear Advice"),
        gr.Textbox(label="Language Info"),
        gr.File(label="Download PDF Report")
    ],
    title="🍲 FoodHealth Pro+",
    description="Upload a food image and get nutrition details, health risk, personalized advice, audio guidance, and a downloadable PDF report."
)

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