import gradio as gr
import pandas as pd
import faiss
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from sentence_transformers import SentenceTransformer

# -------------------------------
# Load dataset
# -------------------------------
file_path = "marketing-campaigns.csv"  # file uploaded in Hugging Face space
df = pd.read_csv(file_path)
df = df.dropna()  # drop completely empty rows
df["text"] = df.astype(str).agg(" | ".join, axis=1)  # merge all cols into text

# -------------------------------
# Embeddings + FAISS
# -------------------------------
embed_model = SentenceTransformer("all-MiniLM-L6-v2")
embeddings = embed_model.encode(df["text"].tolist(), convert_to_tensor=True, show_progress_bar=True)
embeddings_np = embeddings.detach().cpu().numpy()

d = embeddings_np.shape[1]
index = faiss.IndexFlatL2(d)
index.add(embeddings_np)

# -------------------------------
# Load LLM (default Phi-4-mini)
# -------------------------------
model_choices = {
    "openai/gpt-oss-120b": "microsoft/phi-2",
    "openai/gpt-oss-20b": "microsoft/phi-2"  # placeholder, can map to another HF model
}

current_model_id = "openai/gpt-oss-120b"
hf_model = model_choices[current_model_id]

tokenizer = AutoTokenizer.from_pretrained(hf_model)
model = AutoModelForCausalLM.from_pretrained(
    hf_model, torch_dtype=torch.float32, device_map="auto"
)

# -------------------------------
# RAG Functions
# -------------------------------
def retrieve_context(query, k=3):
    query_vec = embed_model.encode([query], convert_to_tensor=True).cpu().numpy()
    D, I = index.search(query_vec, k)
    results = [df.iloc[i]["text"] for i in I[0]]
    return results

def generate_with_rag(prompt, temperature=0.8, max_tokens=250):
    # Step 1: Retrieve context
    context = retrieve_context(prompt, k=3)
    context_str = "\n".join(context)

    # Step 2: Construct grounded prompt
    rag_prompt = f"""
    You are a creative AI campaign assistant.
    Use the following supporting data to ground your answer:
    {context_str}

    Task: Generate a unique and creative marketing campaign idea for:
    {prompt}
    """

    # Step 3: Generate
    inputs = tokenizer(rag_prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_length=max_tokens,
        temperature=temperature,
        top_p=0.9
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

def switch_model(model_choice):
    """Switch between available models."""
    global model, tokenizer, current_model_id
    hf_model = model_choices.get(model_choice, "microsoft/phi-4-mini")
    tokenizer = AutoTokenizer.from_pretrained(hf_model)
    model = AutoModelForCausalLM.from_pretrained(
        hf_model, torch_dtype=torch.float32, device_map="auto"
    )
    current_model_id = model_choice
    return gr.update(visible=(model_choice=="openai/gpt-oss-120b")), gr.update(visible=(model_choice=="openai/gpt-oss-20b")), model_choice

# -------------------------------
# Custom CSS
# -------------------------------
custom_css = """
.gradio-container {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 25%, #f093fb 50%, #4facfe 75%, #00f2fe 100%);
    background-size: 400% 400%;
    animation: gradient-animation 15s ease infinite;
    min-height: 100vh;
}
@keyframes gradient-animation {
    0% { background-position: 0% 50%; }
    50% { background-position: 100% 50%; }
    100% { background-position: 0% 50%; }
}
.main-container {
    background-color: rgba(255, 255, 255, 0.95);
    backdrop-filter: blur(10px);
    border-radius: 20px;
    padding: 20px;
    box-shadow: 0 8px 32px 0 rgba(31, 38, 135, 0.37);
    border: 1px solid rgba(255, 255, 255, 0.18);
    margin: 10px;
}
"""

# -------------------------------
# Gradio UI
# -------------------------------
with gr.Blocks(fill_height=True, theme=gr.themes.Soft, css=custom_css) as demo:
    # Top badges
    with gr.Row(elem_classes="badge-container"):
        gr.HTML("""
            <div style="display: flex; justify-content: center; align-items: center; gap: 15px;">
                <a href="https://huggingface.co/spaces/VIDraft/gpt-oss-RAG" target="_blank">
                    <img src="https://img.shields.io/static/v1?label=gpt-oss-20b&message=RAG&color=%23000080&labelColor=%23ffa500&logo=huggingface&logoColor=white&style=for-the-badge" alt="badge">
                </a>
                <a href="https://discord.gg/openfreeai" target="_blank">
                    <img src="https://img.shields.io/static/v1?label=Discord&message=Openfree%20AI&color=%23000080&labelColor=%23ffa500&logo=discord&logoColor=white&style=for-the-badge" alt="badge">
                </a>
            </div>
        """)

    with gr.Row():
        # Sidebar
        with gr.Column(scale=1):
            with gr.Group(elem_classes="main-container"):
                gr.Markdown("# 🚀 Inference Provider")
                model_dropdown = gr.Dropdown(
                    choices=list(model_choices.keys()),
                    value=current_model_id,
                    label="📊 Select Model"
                )
                reload_btn = gr.Button("🔄 Apply Model Change", variant="primary")

                with gr.Accordion("⚙️ Advanced Options", open=False):
                    temperature = gr.Slider(
                        minimum=0, maximum=2, value=0.8, step=0.1, label="Temperature"
                    )
                    max_tokens = gr.Slider(
                        minimum=50, maximum=1024, value=250, step=10, label="Max Tokens"
                    )

        # Main chat area
        with gr.Column(scale=3):
            with gr.Group(elem_classes="main-container"):
                gr.Markdown("## 💬 RAG-powered Creative Campaign Assistant")

                query = gr.Textbox(label="Enter campaign idea / theme", lines=2)
                output = gr.Textbox(label="Generated Campaign Script", lines=10)

                btn = gr.Button("✨ Generate with RAG")

                btn.click(
                    generate_with_rag,
                    inputs=[query, temperature, max_tokens],
                    outputs=output
                )

                with gr.Column(visible=True) as model_120b_container:
                    gr.Markdown("### Active Model: openai/gpt-oss-120b")
                with gr.Column(visible=False) as model_20b_container:
                    gr.Markdown("### Active Model: openai/gpt-oss-20b")

    reload_btn.click(
        fn=switch_model,
        inputs=[model_dropdown],
        outputs=[model_120b_container, model_20b_container, gr.State(current_model_id)]
    )

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