import gradio as gr
import numpy as np
import pandas as pd
import faiss
import torch
from transformers import AutoTokenizer, AutoModel, pipeline, AutoModelForCausalLM
from huggingface_hub import login
import os
from datetime import datetime

# Authenticate Hugging Face token securely
login(token=os.getenv("HF_TOKEN"))

# Load FAISS index and product data
index = faiss.read_index("deberta_faiss.index")
text_data = pd.read_csv("deberta_text_data.csv")["Retrieved Text"].tolist()

# Load DeBERTa for embedding
deberta_model_name = "microsoft/deberta-v3-base"
deberta_tokenizer = AutoTokenizer.from_pretrained(deberta_model_name)
deberta_model = AutoModel.from_pretrained(deberta_model_name).to("cpu")

# Load Falcon 1B for text generation
llm_model_name = "tiiuae/falcon-rw-1b"
llm_tokenizer = AutoTokenizer.from_pretrained(llm_model_name)
llm_model = AutoModelForCausalLM.from_pretrained(llm_model_name)
llm_pipeline = pipeline("text-generation", model=llm_model, tokenizer=llm_tokenizer, device=-1)

log_path = "query_logs.txt"

def generate_embeddings(queries):
    tokens = deberta_tokenizer(queries, return_tensors="pt", padding=True, truncation=True).to("cpu")
    with torch.no_grad():
        embeddings = deberta_model(**tokens).last_hidden_state.mean(dim=1).cpu().numpy().astype("float32")
    return embeddings

def generate_response(user_query):
    with open(log_path, "a") as log_file:
        log_file.write(f"{datetime.now()} - {user_query}\n")

    query_embedding = generate_embeddings([user_query])
    faiss.normalize_L2(query_embedding)
    distances, indices = index.search(query_embedding, k=5)
    retrieved_docs = [text_data[i] for i in indices[0]]
    context = ", ".join(set(retrieved_docs))

    prompt = f"""
    Using the following product descriptions:
    {context}

    Carefully craft a well-structured response to the following question:

    **Question:** {user_query}

    **Instructions:**
    1. Incorporate **all** retrieved product descriptions.
    2. Use a **formal yet engaging** tone.
    3. Provide **practical & creative** luxury decor ideas.
    4. Ensure a **cohesive & detailed response.**

    **Your response:**
    """

    result = llm_pipeline(prompt, max_new_tokens=300, do_sample=True, truncation=True, pad_token_id=llm_tokenizer.eos_token_id)[0]["generated_text"]
    return result

# Gradio UI with branding
with gr.Blocks(css="""
#logo-img {
    display: block;
    margin-left: auto;
    margin-right: auto;
    margin-top: 5px;
    margin-bottom: 5px;
    border-radius: 12px;
    box-shadow: 0px 2px 10px rgba(0, 0, 0, 0.05);
}
""", theme=gr.themes.Soft()) as demo:

    
    with gr.Column():
        with gr.Row():
            gr.Image(
    value="OnlineRetail.png",
    show_label=False,
    show_download_button=False,
    height=100,
    width=100,
    elem_id="logo-img"
)
        gr.Markdown("### 🛋️ Luxury Decor Assistant")
        gr.Markdown("Ask your decor questions based on real product descriptions. <br> Powered by **DeBERTa + FAISS + Falcon 1B**.")

        with gr.Row():
            input_box = gr.Textbox(label="Textbox", placeholder="e.g. Suggest cozy decor for a small bedroom")
            output_box = gr.Textbox(label="Response")

        submit_btn = gr.Button("✨ Generate Suggestion")
        submit_btn.click(fn=generate_response, inputs=input_box, outputs=output_box)

demo.launch()