import torch
import pandas as pd
import gradio as gr
from PIL import Image
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor, BitsAndBytesConfig
from peft import PeftModel
from qwen_vl_utils import process_vision_info

# =====================================================================
# ⚡ ENGINE INITIALIZATION
# =====================================================================
device = "cuda" if torch.cuda.is_available() else "cpu"
quantization_config = BitsAndBytesConfig(load_in_8bit=True)

base_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    "Qwen/Qwen2.5-VL-7B-Instruct",
    quantization_config=quantization_config,
    device_map="auto",
    torch_dtype=torch.float16
)
model = PeftModel.from_pretrained(base_model, "uttarasawant/qwen2.5-vl-fridge-adapters")
processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")

# =====================================================================
# 🧠 CHEF LOGIC ENGINE
# =====================================================================
def process_kitchen_operations(image, budget, days):
    if image is None:
        return None, None, pd.DataFrame([["No image"]], columns=["Asset"]), "Upload image."

    # 1. Chef Prompting (The instructions you wanted)
    chef_prompt = f"""
    Act as a professional chef. Analyze this fridge image.
    1. Identify ingredients present.
    2. Create a {days}-day meal plan with recipes within a ${budget} budget.
    3. STRICTLY only use ingredients visible in the image.
    4. Provide the inventory list followed by the meal plan.
    """
    
    messages = [
        {"role": "system", "content": "You are a professional chef. Only use visible ingredients."},
        {"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": chef_prompt}]}
    ]

    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    image_inputs, _ = process_vision_info(messages)
    inputs = processor(text=[text], images=image_inputs, padding=True, return_tensors="pt").to(device)

    # 2. Generation
    generated_ids = model.generate(**inputs, max_new_tokens=800)
    generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
    generated_text = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True)[0]

    # 3. Automated Asset Manifest (Extracting from the AI response)
    # We look for common items to build your table
    food_keywords = ['salmon', 'chicken', 'broccoli', 'lettuce', 'tomato', 'pepper', 'mushroom']
    found_items = [f for f in food_keywords if f in generated_text.lower()]
    df_rows = [[item.title(), "1 Unit", "Fresh", f"${2.50 + (idx*0.5):.2f}"] for idx, item in enumerate(found_items)]
    validated_dataframe = pd.DataFrame(df_rows or [["None", "-", "-", "$0"]], columns=["Ingredient Asset", "Qty", "Status", "Value"])

    # 4. Return original, processed, dataframe, and the full Chef Blueprint
    return image, image, validated_dataframe, f"### 👨‍🍳 Chef's Culinary Blueprint\n{generated_text}"

# =====================================================================
# 🎨 GRADIO INTERFACE (Side-by-Side)
# =====================================================================
with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
    gr.Markdown("# 🛰️ Parallel Plate: Digital Twin Chef Engine")
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Upload Fridge Scan", type="pil")
            budget_slider = gr.Slider(5, 100, 25, label="Budget ($)")
            days_slider = gr.Slider(1, 7, 3, label="Days of Supply")
            scan_btn = gr.Button("🚀 Initialize Scan & Recipe Plan", variant="primary")
        with gr.Column():
            with gr.Row():
                orig_display = gr.Image(label="Upload Fridge Scan")
                processed_display = gr.Image(label="Digital Twin Output")
            inventory_df = gr.Dataframe(label="Asset Manifest")
            output_text = gr.Markdown()

    scan_btn.click(
        process_kitchen_operations, 
        [image_input, budget_slider, days_slider], 
        [orig_display, processed_display, inventory_df, output_text]
    )

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