import streamlit as st
import pandas as pd
import chromadb
from sentence_transformers import SentenceTransformer
from transformers import pipeline, AutoModelForQuestionAnswering, AutoTokenizer
from PIL import Image
from io import BytesIO
import requests

# --- 1. Load Recipes Dataset ---
@st.cache_data
def load_recipes():
    try:
        recipes_df = pd.read_csv("recipes.csv")
        recipes_df = recipes_df.rename(columns={"recipe_name": "title", "directions": "instructions"})
        recipes_df = recipes_df[['title', 'ingredients', 'instructions', 'img_src']]
        recipes_df.fillna("", inplace=True)
        recipes_df["ingredients"] = recipes_df["ingredients"].str.lower().str.replace(r'[^\w\s]', '', regex=True)
        recipes_df["combined_text"] = recipes_df["title"] + " " + recipes_df["ingredients"]
        return recipes_df
    except Exception as e:
        st.error(f"⚠ Error loading recipes: {e}")
        return pd.DataFrame()

recipes_df = load_recipes()

# --- 2. Load SentenceTransformer Model ---
@st.cache_resource
def load_embedding_model():
    return SentenceTransformer("all-mpnet-base-v2")

embedding_model = load_embedding_model()

# --- 3. Initialize ChromaDB ---
chroma_client = chromadb.PersistentClient(path="./chroma_db")
collection = chroma_client.get_or_create_collection(name="recipe_collection")

# --- 4. Generate & Store Embeddings ---
def get_sentence_transformer_embeddings(text):
    return embedding_model.encode(text).tolist()

try:
    existing_data = collection.get()
    existing_ids = set(existing_data["ids"]) if existing_data and "ids" in existing_data else set()
except Exception as e:
    st.error(f"⚠ ChromaDB Error: {e}")
    existing_ids = set()

for index, row in recipes_df.iterrows():
    recipe_id = str(index)
    if recipe_id in existing_ids:
        continue
    embedding = get_sentence_transformer_embeddings(row["combined_text"])
    if embedding:
        collection.add(embeddings=[embedding], documents=[row["combined_text"]], ids=[recipe_id])

# --- 5. Retrieve Similar Recipes ---
def retrieve_recipes(query, top_k=3):
    query_embedding = get_sentence_transformer_embeddings(query)
    results = collection.query(query_embeddings=[query_embedding], n_results=top_k)
    
    if results and "ids" in results and results["ids"]:
        recipe_indices = [int(id) for id in results["ids"][0] if id.isdigit()]
        return recipes_df.iloc[recipe_indices] if recipe_indices else None
    return None

# --- 6. Load LLM Model ---
@st.cache_resource
def load_llm_model():
    tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")
    model = AutoModelForQuestionAnswering.from_pretrained("deepset/roberta-base-squad2")
    return pipeline("question-answering", model=model, tokenizer=tokenizer)

llm_model = load_llm_model()

# --- 7. Load Classifier ---
@st.cache_resource
def load_classifier():
    return pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

classifier = load_classifier()

# --- 8. Query Classification ---
def is_food_related(query):
    labels = ["Food & Cooking", "General Knowledge"]
    result = classifier(query, candidate_labels=labels, multi_label=False)
    return result["labels"][0] == "Food & Cooking"

def answer_question(query, context=""):
    greetings = ["hi", "hello", "hii", "hey", "greetings", "how are you", "what's up", "how's it going"]
    query_cleaned = query.lower().strip()
    
    if query_cleaned in greetings:
        return "Hello! I'm here to assist with recipes and food-related questions. 🍽️ What would you like to know?"
    
    if not is_food_related(query_cleaned):
        return "I specialize in food-related topics! 🍽️ Please ask me about recipes, cooking methods, or ingredients."
    
    # Retrieve related recipes only if the query is about food
    related_recipes = retrieve_recipes(query, top_k=1)
    if related_recipes is None or related_recipes.empty:
        return "I specialize in recipes! 🍽️ Feel free to ask me about ingredients, cooking methods, or meal ideas. 😊"
    
    # Use the most relevant recipe's instructions for answer generation
    context = related_recipes.iloc[0]['instructions']
    response = llm_model(question=query, context=context)
    answer = response.get("answer", "").strip()
    
    return answer if answer else "I'm not sure, but I can help with recipes! What would you like to cook?"

# --- 9. Display Image Function ---
def display_image(image_url, recipe_name):
    try:
        if not isinstance(image_url, str) or not image_url.startswith("http"):
            raise ValueError("Invalid or missing image URL")
        response = requests.get(image_url, timeout=5)
        response.raise_for_status()
        image = Image.open(BytesIO(response.content))
        st.image(image, caption=recipe_name, use_container_width=True)
    except requests.exceptions.RequestException as e:
        st.warning(f"⚠ Image fetch error: {e}")
        placeholder_url = "https://via.placeholder.com/300?text=No+Image"
        st.image(placeholder_url, caption=recipe_name, use_container_width=True)

# --- 10. Streamlit UI ---
st.title("🍽️ AI Recipe & Q&A Assistant")

user_query = st.text_input("Enter your question or recipe search query:", "", key="main_query_input")

if "retrieved_recipes" not in st.session_state:
    st.session_state["retrieved_recipes"] = None

if st.button("Ask AI"):
    if user_query:
        if is_food_related(user_query):  # First check if it's a food-related query
            response = answer_question(user_query)
            st.subheader("🤖 AI Answer:")
            st.write(response)

            retrieved_recipes = retrieve_recipes(user_query)
            if retrieved_recipes is not None and not retrieved_recipes.empty:
                st.session_state["retrieved_recipes"] = retrieved_recipes
                st.subheader("🍴 Found Recipes:")
                for _, recipe in retrieved_recipes.iterrows():
                    st.markdown(f"### {recipe['title']}")
                    st.write(f"**Ingredients:** {recipe['ingredients']}")
                    st.write(f"**Instructions:** {recipe['instructions']}")
                    display_image(recipe.get('img_src', ''), recipe['title'])
            else:
                st.warning("⚠️ No relevant recipes found.")
        else:
            # If query is not food-related, return only general response (No recipe retrieval)
            st.subheader("🤖 AI Answer:")
            st.write(answer_question(user_query))