import os
import gradio as gr
from pinecone import Pinecone, ServerlessSpec
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext, Settings
from llama_index.vector_stores.pinecone import PineconeVectorStore
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI

# --- System Prompt ---
SYSTEM_PROMPT = """You are Shifu, a polite and professional Healthy Recipe assistant.
Answer ONLY using the information found in the indexed recipe document(s).
If the answer is not in the document(s), say: "I couldn’t find that in the document."
Keep responses concise, helpful, and friendly.
"""

# ===== CONFIG =====
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
OPENAI_API_KEY   = os.getenv("OPENAI_API_KEY")
if not PINECONE_API_KEY or not OPENAI_API_KEY:
    raise RuntimeError("Missing PINECONE_API_KEY or OPENAI_API_KEY")

DATA_DIR = "data"
BANNER_PATH = os.path.join(DATA_DIR, "shafaq_banner.png")
if not os.path.exists(BANNER_PATH):
    raise RuntimeError("Banner not found at: data/shafaq_banner.png")

EMBED_MODEL = "text-embedding-3-small"
LLM_MODEL   = "gpt-4o-mini"
TOP_K       = 4

# ===== LlamaIndex & Pinecone =====
Settings.embed_model = OpenAIEmbedding(model=EMBED_MODEL, api_key=OPENAI_API_KEY)
Settings.llm = OpenAI(model=LLM_MODEL, api_key=OPENAI_API_KEY, system_prompt=SYSTEM_PROMPT)

pc = Pinecone(api_key=PINECONE_API_KEY)

def ensure_index(name: str, dim: int = 1536):
    names = [i["name"] for i in pc.list_indexes()]
    if name not in names:
        pc.create_index(
            name=name, dimension=dim, metric="cosine",
            spec=ServerlessSpec(cloud="aws", region="us-east-1"),
        )
    return pc.Index(name)

pinecone_index = ensure_index("shafaq-recipes-index", dim=1536)
vector_store = PineconeVectorStore(pinecone_index=pinecone_index)

def bootstrap_index():
    if not os.path.isdir(DATA_DIR):
        raise RuntimeError("No 'data/' directory found.")
    docs = SimpleDirectoryReader(DATA_DIR).load_data()
    if not docs:
        raise RuntimeError("No documents found in data/.")
    storage_ctx = StorageContext.from_defaults(vector_store=vector_store)
    VectorStoreIndex.from_documents(docs, storage_context=storage_ctx, show_progress=True)

bootstrap_index()

def answer(query: str) -> str:
    if not query.strip():
        return "Please enter a question."
    index = VectorStoreIndex.from_vector_store(vector_store)
    resp = index.as_query_engine(similarity_top_k=TOP_K).query(query)
    return str(resp)

FAQS = [
    "",
    "How long do these breakfast recipes last in the fridge?",
    "Can I freeze any of these recipes?",
    "Are there gluten-free breakfast options?",
    "How can I make these breakfasts dairy-free?",
    "What can I use instead of eggs?",
    "Can I use different veggies or toppings?",
    "I don’t have almond or coconut milk. What else can I use?",
    "How can I add more flavor to these recipes?",
    "What if I don’t have protein powder?",
    "Can I prep these meals ahead of time?",
]

def use_faq(selected_faq: str, free_text: str):
    prompt = (selected_faq or "").strip() or (free_text or "").strip()
    if not prompt:
        return "", "Please select a FAQ or type your question."
    return prompt, answer(prompt)

# ===== UI =====
CSS = """
.header { text-align: left; }
.logo img { width:100%; max-width: 320px; border-radius: 12px; }
.title { font-weight: 700; font-size: 1.8rem; color: #41644A; margin-bottom: 4px; }
.subnote { font-size: 1rem; color: #6C757D; margin-bottom: 20px; }
.section-title { font-weight: 600; margin-top: 12px; }
"""

with gr.Blocks(css=CSS, theme=gr.themes.Soft(primary_hue="green", secondary_hue="orange")) as demo:
    with gr.Row():
        with gr.Column(scale=1, min_width=350):
            gr.Image(value=BANNER_PATH, show_label=False, elem_classes=["logo"])

        with gr.Column(scale=2):
            gr.Markdown("""
            <div style="text-align: left; margin-bottom: 1.5rem;">
                <h1 style="font-size: 2.4rem; font-weight: 800; color: #2F4F4F; margin: 0;">
                    Shafaq’s Breakfast Buddy
                </h1>
                <p style="font-size: 1.2rem; color: #6C757D; margin-top: 6px;">
                    Healthy Recipe Recommender based on what’s in your kitchen 🍳🥑
                </p>
            </div>
            """)

            gr.Markdown("#### Ask from Frequently Asked Questions")
            faq = gr.Dropdown(choices=FAQS, value=FAQS[0], label="Select a common question")

            gr.Markdown("#### Or type your question")
            user_q = gr.Textbox(
                label="Your question",
                placeholder="e.g., Suggest a healthy breakfast with oats and banana",
                lines=2
            )

            ask_btn = gr.Button("Get Recipe 💡", variant="primary")

            chosen_prompt = gr.Textbox(label="Query sent", interactive=False)
            answer_box = gr.Markdown()

            # ✅ Must be INSIDE the same column/block
            ask_btn.click(use_faq, inputs=[faq, user_q], outputs=[chosen_prompt, answer_box])



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