import os
from dotenv import load_dotenv
import gradio as gr
from string import Template
from typing import List
from pydantic import BaseModel, Field

# ----- Load API Keys -----
load_dotenv()
google_api_key = os.getenv('GOOGLE_API_KEY')
tavily_api_key = os.getenv('TAVILY_API_KEY')

# ----- Pydantic Models -----
class NutritionInfo(BaseModel):
    calories: str = Field(..., description="Total calories in the meal")
    protein: str = Field(..., description="Protein content in grams")
    carbs: str = Field(..., description="Carbohydrates content in grams")
    fat: str = Field(..., description="Fat content in grams")

class BudgetMeal(BaseModel):
    meal_name: str = Field(..., description="Name of the meal")
    ingredients: List[str] = Field(..., description="List of ingredients with quantities")
    cooking_steps: List[str] = Field(..., description="Step-by-step cooking instructions")
    nutrition_info: NutritionInfo = Field(..., description="Nutritional breakdown of the meal")
    reason_for_selection: str = Field(..., description="Explanation for why this meal was chosen")

class BudgetBasedMeal(BaseModel):
    low_budget: BudgetMeal = Field(..., description="Low-budget version of the meal")
    medium_budget: BudgetMeal = Field(..., description="Medium-budget version of the meal")
    high_budget: BudgetMeal = Field(..., description="High-budget version of the meal")

class ThreeMealPlan(BaseModel):
    meal_type: str
    breakfast: BudgetBasedMeal
    lunch: BudgetBasedMeal
    dinner: BudgetBasedMeal

class FourMealPlan(BaseModel):
    meal_type: str
    breakfast: BudgetBasedMeal
    lunch: BudgetBasedMeal
    snack: BudgetBasedMeal
    dinner: BudgetBasedMeal

class IntermittentFastingPlan(BaseModel):
    meal_type: str
    lunch: BudgetBasedMeal
    dinner: BudgetBasedMeal

# ----- Prompt Template -----
prompt_template = Template(
    """
You are a top nutritionist specializing in personalized meal planning. Based on the user's profile, create a personalized one-day meal plan that STRICTLY adheres to their dietary preferences and COMPLETELY EXCLUDES any ingredients they are allergic to. The user follows the "$meal_plan_type" pattern.

IMPORTANT DIETARY GUIDELINES:
1. STRICTLY follow the user's dietary preference: $dietary_preference
2. ABSOLUTELY AVOID any ingredients listed in allergies/restrictions: $allergies
   - Double-check each ingredient to ensure NO allergens are included
   - If an ingredient could contain hidden allergens, choose a safe alternative

Meal Plan Type Guide:
- "3 meals/day" → breakfast, lunch, dinner
- "4 meals/day" → breakfast, lunch, snack, dinner
- "Intermittent fasting (2 meals)" → lunch, dinner

Each meal should have **3 flexible options** while maintaining dietary requirements:
- low budget (affordable while meeting dietary needs)
- medium budget (balanced options within dietary restrictions)
- high budget (premium ingredients following dietary preferences)

Each option must include:
- Meal name (clearly indicating it follows dietary preferences)
- Ingredients (all safe and compliant with dietary restrictions)
- Cooking steps
- Basic nutrition info (calories, protein, carbs, fat)
- Short reason for choosing this meal based on the user's chosen package and dietary needs

User Profile:
- Age group: $age_group
- Height: $height inches
- Weight: $weight lbs
- Gender: $gender
- Dietary preference: $dietary_preference (STRICT ADHERENCE REQUIRED)
- Allergies or restrictions: $allergies (MUST BE COMPLETELY AVOIDED)
- Goal/package: $package

Output format must be clean and JSON-like, without extra keys. Just the meals as per plan type with 3 budget-based options each. Every meal MUST comply with dietary preferences and exclude allergens.
"""
)

# ----- Tavily Tool -----
from langchain_tavily import TavilySearch
tavily_tool = TavilySearch(
    max_results=20,
    topic="general",
    include_answer=True,
    include_raw_content=True,
    search_depth="advanced",
    tavily_api_key=tavily_api_key,
    include_domains=[
        "https://www.nutritionvalue.org/",
        "https://www.walmart.com/search?q=",
        "https://www.healthline.com/nutrition",
        "https://www.healthline.com/nutrition/meal-kits",
        "https://www.healthline.com/nutrition/meal-kits/diets",
        "https://www.healthline.com/nutrition/special-diets",
        "https://www.healthline.com/nutrition/healthy-eating",
        "https://www.healthline.com/nutrition/food-freedom",
        "https://www.healthline.com/nutrition/feel-good-food",
        "https://www.healthline.com/nutrition/products",
        "https://www.healthline.com/nutrition/vitamins-supplements",
        "https://www.healthline.com/nutrition/sustain",
    ],
)

# ----- LLM + Agents -----
from langchain_google_genai import ChatGoogleGenerativeAI
llm = ChatGoogleGenerativeAI(model="gemini-2.5-pro-preview-03-25", google_api_key=google_api_key)

from langgraph.prebuilt import create_react_agent
agent_3_meals = create_react_agent(
    llm,
    tools=[tavily_tool],
    response_format=ThreeMealPlan
)

agent_4_meals = create_react_agent(
    llm,
    tools=[tavily_tool],
    response_format=FourMealPlan
)

agent_intermittent = create_react_agent(
    llm,
    tools=[tavily_tool],
    response_format=IntermittentFastingPlan
)

# ----- Render Functions -----
def render_meal(meal: BudgetMeal, budget_label: str) -> str:
    ingredients_str = "\n- ".join(meal.ingredients)
    steps_str = "\n1. ".join(meal.cooking_steps)
    return (
        f"### {budget_label} Option\n\n"
        f"**🍽️ Meal Name**: {meal.meal_name}\n\n"
        f"**📝 Ingredients**:\n- {ingredients_str}\n\n"
        f"**👨‍🍳 Cooking Steps**:\n1. {steps_str}\n\n"
        f"**🍎 Nutrition Info**:\n"
        f"- Calories: {meal.nutrition_info.calories}\n"
        f"- Protein: {meal.nutrition_info.protein}\n"
        f"- Carbs: {meal.nutrition_info.carbs}\n"
        f"- Fat: {meal.nutrition_info.fat}\n\n"
        f"**💡 Why this meal?**\n{meal.reason_for_selection}\n"
    )

def render_budget_meal(meal_obj: BudgetBasedMeal, meal_type: str) -> str:
    return (
        f"## 🍱 {meal_type.title()}\n\n"
        f"{render_meal(meal_obj.low_budget, 'Low Budget')}\n"
        f"{render_meal(meal_obj.medium_budget, 'Medium Budget')}\n"
        f"{render_meal(meal_obj.high_budget, 'High Budget')}\n"
    )

def format_three_meal_plan(plan: ThreeMealPlan) -> str:
    return (
        f"# 🧾 Meal Plan: {plan.meal_type}\n\n"
        f"{render_budget_meal(plan.breakfast, 'Breakfast')}\n"
        f"{render_budget_meal(plan.lunch, 'Lunch')}\n"
        f"{render_budget_meal(plan.dinner, 'Dinner')}\n"
    )

def format_four_meal_plan(plan: FourMealPlan) -> str:
    return (
        f"# 🧾 Meal Plan: {plan.meal_type}\n\n"
        f"{render_budget_meal(plan.breakfast, 'Breakfast')}\n"
        f"{render_budget_meal(plan.lunch, 'Lunch')}\n"
        f"{render_budget_meal(plan.snack, 'Snack')}\n"
        f"{render_budget_meal(plan.dinner, 'Dinner')}\n"
    )

def format_if_plan(plan: IntermittentFastingPlan) -> str:
    return (
        f"# 🧾 Meal Plan: {plan.meal_type}\n\n"
        f"{render_budget_meal(plan.lunch, 'Lunch')}\n"
        f"{render_budget_meal(plan.dinner, 'Dinner')}\n"
    )

# ----- Generate Meal Plan Function -----
def generate_meal_plan(age_group, feet, inches, weight, gender, meal_plan_type, dietary_preference, allergies, package):
    # Compute total height in inches:
    total_height = int(feet) * 12 + int(inches)
    user_input = {
        "age_group": age_group,
        "height": str(total_height),
        "weight": str(weight),
        "gender": gender,
        "meal_plan_type": meal_plan_type,
        "dietary_preference": dietary_preference,
        "allergies": allergies,
        "package": package
    }
    filled_prompt = prompt_template.substitute(**user_input)
    inputs = {"messages": [("user", filled_prompt)]}
    
    if meal_plan_type.startswith("3 meals/day"):
        result = agent_3_meals.invoke(inputs)["structured_response"]
        formatted = format_three_meal_plan(result)
    elif meal_plan_type.startswith("4 meals/day"):
        result = agent_4_meals.invoke(inputs)["structured_response"]
        formatted = format_four_meal_plan(result)
    else:  # Intermittent Fasting
        result = agent_intermittent.invoke(inputs)["structured_response"]
        formatted = format_if_plan(result)
    
    return formatted

# ----- Gradio UI -----
demo = gr.Blocks()

with demo:
    gr.Markdown("## 🍽️ Personalized Meal Plan Generator")
    with gr.Row():
        age_group = gr.Dropdown(choices=["18-24", "25-30", "31-40", "41-50", "51+"], label="Age Group")
        gender = gr.Dropdown(choices=["male", "female", "other"], label="Gender")
    with gr.Row():
        feet = gr.Number(label="Height (feet)")
        inches = gr.Number(label="Height (inches)")
        weight = gr.Number(label="Weight (lbs)")
    meal_plan_type = gr.Radio(
        choices=[
            "3 meals/day (Breakfast, Lunch, Dinner)",
            "4 meals/day (Breakfast, Lunch, Snack, Dinner)",
            "Intermittent fasting (2 meals)"
        ],
        label="Meal Plan Type"
    )
    dietary_preference = gr.Dropdown(
        choices=["Keto", "Vegan", "Vegetarian", "Low-Carb", "High-Protein", "Balanced"],
        label="Dietary Preference"
    )
    allergies = gr.Textbox(label="Allergies or Restrictions (e.g., Gluten, Dairy, Nuts or 'None')")
    package = gr.Dropdown(
        choices=["Fitness and Mobility", "Focus Flow", "No More Insomnia"],
        label="Goal/Package"
    )
    
    with gr.Row():
        # Add a status indicator
        status_indicator = gr.Markdown("Status: Ready")
    
    generate_btn = gr.Button("Generate Meal Plan")
    output_display = gr.Markdown()
    
    # Add the loading indicator logic
    def generate_with_loading(age_group, feet, inches, weight, gender, meal_plan_type, dietary_preference, allergies, package):
        # Return a loading message for the status indicator
        return "Status: Generating your meal plan... This may take a moment! ⏳"
    
    def finalize_generation(age_group, feet, inches, weight, gender, meal_plan_type, dietary_preference, allergies, package):
        # Generate the meal plan
        result = generate_meal_plan(age_group, feet, inches, weight, gender, meal_plan_type, dietary_preference, allergies, package)
        # Update the status indicator
        return "Status: Ready", result
    
    # Connect the button click to both functions in sequence
    generate_btn.click(
        fn=generate_with_loading,
        inputs=[age_group, feet, inches, weight, gender, meal_plan_type, dietary_preference, allergies, package],
        outputs=status_indicator,
    ).then(
        fn=finalize_generation,
        inputs=[age_group, feet, inches, weight, gender, meal_plan_type, dietary_preference, allergies, package],
        outputs=[status_indicator, output_display],
    )

demo.launch()