# =====================================================================
# 🛸 HUGGING FACE ZERO-GPU INITIALIZATION (MUST BE FIRST)
# =====================================================================
import sys
# This forces spaces to load right away if it's installed in the HF container
try:
    import spaces
except ImportError:
    pass
    
import torch
import pandas as pd
import gradio as gr
import cv2
import numpy as np
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")

def sample_frames_from_video(video_path, num_frames=4):
    cap = cv2.VideoCapture(video_path)
    frames = []
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    # Handle empty videos
    if total_frames <= 0: return []
    indices = np.linspace(0, total_frames - 1, num_frames, dtype=int)
    for i in range(total_frames):
        ret, frame = cap.read()
        if not ret: break
        if i in indices:
            frames.append(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))
    cap.release()
    return frames

# =====================================================================
# 🧠 CHEF LOGIC ENGINE
# =====================================================================
@spaces.GPU(duration=120)
def process_kitchen_operations(media_input, budget, days):
    # GUARD: Stop if no input provided
    if media_input is None:
        return None, None, pd.DataFrame(columns=["Ingredient Asset", "Qty", "Status", "Value"]), "### ⚠️ System Idle\nPlease upload an image or video."

    if isinstance(media_input, str): 
        images = sample_frames_from_video(media_input)
        if not images: return None, None, pd.DataFrame(), "### ❌ Error\nCould not extract frames from video."
    else: 
        images = [media_input]

    chef_prompt = f"Act as a professional chef. Identify ingredients. Create a {days}-day meal plan (budget ${budget}). Output as Markdown Table (Day|Breakfast|Lunch|Dinner). Provide inventory list first."
    
    content = [{"type": "image", "image": img} for img in images]
    content.append({"type": "text", "text": chef_prompt})
    
    messages = [{"role": "system", "content": "You are a professional chef. Only use visible ingredients."}, {"role": "user", "content": content}]

    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)

    generated_ids = model.generate(**inputs, max_new_tokens=400)
    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]

    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)]
    df = pd.DataFrame(df_rows or [["None", "-", "-", "$0"]], columns=["Ingredient Asset", "Qty", "Status", "Value"])

    return images[0], images[0], df, f"### 👨‍🍳 Chef's Culinary Blueprint\n{generated_text}"

# =====================================================================
# 🎨 GRADIO INTERFACE
# =====================================================================
with gr.Blocks() as demo:
    gr.Markdown("# 🛰️ Parallel Plate: Digital Twin Chef Engine")
    
    with gr.Tabs():
        with gr.TabItem("Upload Image"):
            img_input = gr.Image(type="pil")
        with gr.TabItem("Upload Video"):
            vid_input = gr.Video()

    # Clear other tab when one is used
    img_input.change(fn=lambda: None, outputs=vid_input)
    vid_input.change(fn=lambda: None, outputs=img_input)
    
    budget_slider = gr.Slider(5, 100, 25, label="Budget ($)")
    days_slider = gr.Slider(1, 7, 3, label="Days of Supply")
    
    with gr.Row():
        scan_btn = gr.Button("🚀 Initialize Scan & Recipe Plan", variant="primary")
        clear_btn = gr.Button("🧹 Clear")
    
    with gr.Row():
        orig_display = gr.Image(label="Input Source")
        processed_display = gr.Image(label="Digital Twin Output")
    inventory_df = gr.Dataframe(label="Asset Manifest")
    output_text = gr.Markdown()

    def clear_interface():
        empty_df = pd.DataFrame(columns=["Ingredient Asset", "Qty", "Status", "Value"])
        return [None, None, None, None, empty_df, ""]

    clear_btn.click(fn=clear_interface, inputs=[], outputs=[img_input, vid_input, orig_display, processed_display, inventory_df, output_text])
    
    # Helper to pick the active input
    def choose_input(img, vid): return vid if vid else img

    scan_btn.click(
        fn=lambda img, vid, b, d: process_kitchen_operations(choose_input(img, vid), b, d),
        inputs=[img_input, vid_input, budget_slider, days_slider],
        outputs=[orig_display, processed_display, inventory_df, output_text]
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Monochrome())