import spaces  # For ZeroGPU decorator
import gradio as gr

from transformers import pipeline
from diffusers import StableDiffusionPipeline
import torch
import requests
import re
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
from tempfile import NamedTemporaryFile
from gtts import gTTS

# --- USDA API Setup ---
USDA_API_KEY = "gcwe1gXmMveg7buqddggl6wAZa7Sd7wrZV87P31z"
USDA_SEARCH_URL = "https://api.nal.usda.gov/fdc/v1/foods/search"

def get_nutrients(ingredient):
    params = {"api_key": USDA_API_KEY, "query": ingredient, "pageSize": 1}
    try:
        res = requests.get(USDA_SEARCH_URL, params=params)
        data = res.json()
        if 'foods' in data and len(data['foods']) > 0:
            return {
                n['nutrientName']: n['value'] for n in data['foods'][0].get('foodNutrients', [])
                if n.get('nutrientName')
            }
        return {}
    except:
        return {}

def get_calories(ingredient):
    nutrients = get_nutrients(ingredient)
    return nutrients.get("Energy", 0.0)

# --- Load CPU-only model globally (safe) ---
recipe_model = pipeline("text-generation", model="samdak93/qrit-2")

# Globals to cache GPU models after loading
image_model = None
recipe_model_gpu = None  # If you want to move recipe_model to GPU, handle similarly

# --- Utilities ---
forbidden = ["pork", "bacon", "ham", "lard", "gelatin", "alcohol", "beer", "wine", "rum", "whiskey", "vodka", "gin"]

def parse_ingredients(text):
    match = re.search(r"(?i)ingredients:(.*?)(\n|directions:|instructions:|$)", text, re.DOTALL)
    if match:
        ingredients_block = match.group(1).strip()
        return [line.strip("- *") for line in ingredients_block.split("\n") if line.strip()]
    return []

def generate_recipe(key_ingredient):
    for _ in range(5):
        prompt = f"Create a halal recipe under 2000 calories that includes {key_ingredient}.\nIngredients:"
        out = recipe_model(prompt, max_length=300, num_return_sequences=1)[0]['generated_text']
        if any(h in out.lower() for h in forbidden):
            continue

        ingredients = parse_ingredients(out)
        total_cal = sum(get_calories(i) for i in ingredients)
        if total_cal <= 2000 and total_cal > 0:
            return out, ingredients, total_cal
    return "No valid recipe.", [], 0

def get_nutrient_breakdown(ingredients):
    breakdown = {}
    for ing in ingredients:
        breakdown[ing] = get_nutrients(ing)
    return breakdown

@spaces.GPU(duration=120)
def load_image_model():
    global image_model
    if image_model is None:
        image_model = StableDiffusionPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5",
            torch_dtype=torch.float16
        ).to("cuda")
    return image_model

@spaces.GPU(duration=120)
def generate_image(prompt):
    pipe = load_image_model()
    with torch.autocast("cuda"):
        image = pipe(prompt).images[0]
    return image

def export_pdf(recipe_text, nutrients):
    temp_pdf = NamedTemporaryFile(delete=False, suffix=".pdf")
    c = canvas.Canvas(temp_pdf.name, pagesize=letter)
    width, height = letter

    c.setFont("Helvetica", 12)
    c.drawString(50, height - 50, "Generated Halal Recipe")

    text_obj = c.beginText(50, height - 70)
    for line in recipe_text.split('\n'):
        text_obj.textLine(line)
    c.drawText(text_obj)

    y = text_obj.getY() - 20
    c.drawString(50, y, "Nutrient Breakdown:")
    y -= 20
    for ing, data in nutrients.items():
        c.drawString(50, y, f"- {ing}")
        y -= 15
        for k, v in data.items():
            c.drawString(70, y, f"{k}: {v}")
            y -= 12
        y -= 5
        if y < 100:
            c.showPage()
            y = height - 50
    c.save()
    return temp_pdf.name

def generate_audio(recipe_text):
    tts = gTTS(text=recipe_text)
    temp_audio = NamedTemporaryFile(delete=False, suffix=".mp3")
    tts.save(temp_audio.name)
    return temp_audio.name

def app(key_ingredient):
    recipe_text, ingredients, calories = generate_recipe(key_ingredient)
    if not ingredients:
        return "No valid recipe generated.", None, None, None, None

    title_match = re.search(r"(?i)^title: (.+)$", recipe_text, re.MULTILINE)
    title = title_match.group(1).strip() if title_match else f"Dish with {key_ingredient}"

    image = generate_image(f"A delicious {title}, beautifully plated")
    nutrients = get_nutrient_breakdown(ingredients)
    pdf_path = export_pdf(recipe_text, nutrients)
    audio_path = generate_audio(recipe_text)

    display_text = f"{recipe_text}\n\nEstimated Calories: {int(calories)} kcal"
    return display_text, image, nutrients, pdf_path, audio_path

# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("## 🕌 Halal Recipe Generator - Advanced Edition")
    with gr.Row():
        key_input = gr.Textbox(label="Key Ingredient")
        submit_btn = gr.Button("Generate Recipe")

    output_text = gr.Textbox(label="Recipe Text")
    output_img = gr.Image(label="Generated Dish")
    output_table = gr.JSON(label="Nutrient Breakdown")
    output_pdf = gr.File(label="Download Recipe PDF")
    output_audio = gr.Audio(label="Recipe Narration", autoplay=False)

    submit_btn.click(app, inputs=key_input,
                     outputs=[output_text, output_img, output_table, output_pdf, output_audio])

demo.launch(share=True, pwa=True, debug=True)