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

# -------------------------------
# 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 with enterprise styling
# -------------------------------
custom_css = """
:root {
    --primary: #2563eb;
    --primary-dark: #1d4ed8;
    --secondary: #64748b;
    --dark: #1e293b;
    --light: #f8fafc;
    --success: #10b981;
    --warning: #f59e0b;
    --danger: #ef4444;
    --gray-100: #f1f5f9;
    --gray-200: #e2e8f0;
    --gray-300: #cbd5e1;
    --gray-700: #334155;
    --radius: 12px;
    --shadow: 0 10px 15px -3px rgba(0, 0, 0, 0.1), 0 4px 6px -2px rgba(0, 0, 0, 0.05);
}

.gradio-container {
    font-family: 'Inter', 'Segoe UI', system-ui, sans-serif;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    min-height: 100vh;
    padding: 20px;
}

.main-container {
    background-color: rgba(255, 255, 255, 0.95);
    backdrop-filter: blur(10px);
    border-radius: var(--radius);
    padding: 24px;
    box-shadow: var(--shadow);
    border: 1px solid rgba(255, 255, 255, 0.18);
    margin-bottom: 20px;
}

.dark .main-container {
    background-color: rgba(30, 41, 59, 0.95);
    color: var(--light);
}

.header {
    display: flex;
    justify-content: space-between;
    align-items: center;
    padding: 16px 0;
    margin-bottom: 24px;
    border-bottom: 1px solid var(--gray-200);
}

.dark .header {
    border-bottom: 1px solid var(--gray-700);
}

.logo {
    display: flex;
    align-items: center;
    gap: 12px;
    font-weight: 700;
    font-size: 24px;
    color: var(--dark);
}

.dark .logo {
    color: var(--light);
}

.logo-icon {
    background: var(--primary);
    color: white;
    width: 40px;
    height: 40px;
    border-radius: 8px;
    display: flex;
    align-items: center;
    justify-content: center;
}

.nav-links {
    display: flex;
    gap: 24px;
}

.nav-link {
    color: var(--secondary);
    text-decoration: none;
    font-weight: 500;
    transition: color 0.2s;
}

.nav-link:hover, .nav-link.active {
    color: var(--primary);
}

.dark .nav-link {
    color: var(--gray-300);
}

.dark .nav-link:hover, .dark .nav-link.active {
    color: var(--primary);
}

.badge-container {
    margin-bottom: 24px;
}

h1, h2, h3, h4 {
    color: var(--dark);
    font-weight: 700;
}

.dark h1, .dark h2, .dark h3, .dark h4 {
    color: var(--light);
}

.label {
    font-weight: 600;
    color: var(--dark);
    margin-bottom: 8px;
}

.dark .label {
    color: var(--light);
}

input, textarea {
    border-radius: var(--radius) !important;
    padding: 12px 16px !important;
    border: 1px solid var(--gray-300) !important;
    background: white !important;
}

.dark input, .dark textarea {
    border: 1px solid var(--gray-700) !important;
    background: var(--dark) !important;
    color: var(--light) !important;
}

button {
    border-radius: var(--radius) !important;
    padding: 12px 24px !important;
    font-weight: 600 !important;
    transition: all 0.2s !important;
    border: none !important;
}

button.primary {
    background: var(--primary) !important;
    color: white !important;
}

button.primary:hover {
    background: var(--primary-dark) !important;
    transform: translateY(-2px);
}

.dropdown {
    border-radius: var(--radius) !important;
}

.accordion {
    border-radius: var(--radius) !important;
    border: 1px solid var(--gray-200) !important;
}

.dark .accordion {
    border: 1px solid var(--gray-700) !important;
}

.footer {
    text-align: center;
    padding: 24px 0;
    margin-top: 40px;
    color: var(--secondary);
    font-size: 14px;
}

.dark .footer {
    color: var(--gray-300);
}

.stats {
    display: flex;
    gap: 24px;
    margin: 20px 0;
}

.stat-card {
    background: white;
    padding: 16px;
    border-radius: var(--radius);
    box-shadow: var(--shadow);
    flex: 1;
    text-align: center;
}

.dark .stat-card {
    background: var(--dark);
}

.stat-value {
    font-size: 24px;
    font-weight: 700;
    color: var(--primary);
}

.stat-label {
    font-size: 14px;
    color: var(--secondary);
}

.processing-bar {
    height: 4px;
    width: 100%;
    background: var(--gray-200);
    border-radius: 2px;
    overflow: hidden;
    margin: 16px 0;
}

.progress {
    height: 100%;
    background: var(--primary);
    width: 0%;
    transition: width 0.4s;
}

@keyframes pulse {
    0% { opacity: 1; }
    50% { opacity: 0.5; }
    100% { opacity: 1; }
}

.pulse {
    animation: pulse 2s infinite;
}

.model-badge {
    display: inline-block;
    padding: 4px 12px;
    background: var(--gray-100);
    color: var(--gray-700);
    border-radius: 20px;
    font-size: 12px;
    font-weight: 600;
    margin-left: 12px;
}

.dark .model-badge {
    background: var(--gray-700);
    color: var(--gray-200);
}

.tabs {
    display: flex;
    gap: 8px;
    margin-bottom: 20px;
}

.tab {
    padding: 8px 16px;
    background: var(--gray-100);
    border-radius: var(--radius);
    cursor: pointer;
    font-weight: 500;
}

.tab.active {
    background: var(--primary);
    color: white;
}

.dark .tab {
    background: var(--gray-700);
    color: var(--gray-200);
}

.dark .tab.active {
    background: var(--primary);
    color: white;
}
"""
js = """
<script>
// Wait for the page to load
document.addEventListener('DOMContentLoaded', function() {
    // Add header navigation
    const headerHtml = `
        <div class="header">
            <div class="logo">
                <div class="logo-icon">AI</div>
                <span>CampaignForge</span>
            </div>
            <div class="nav-links">
                <a href="#" class="nav-link active">Home</a>
                <a href="#" class="nav-link">Campaigns</a>
                <a href="#" class="nav-link">Analytics</a>
                <a href="#" class="nav-link">Settings</a>
            </div>
        </div>
    `;
    
    // Add stats cards
    const statsHtml = `
        <div class="stats">
            <div class="stat-card">
                <div class="stat-value">${Math.floor(Math.random() * 1000) + 500}</div>
                <div class="stat-label">Campaigns Generated</div>
            </div>
            <div class="stat-card">
                <div class="stat-value">${Math.floor(Math.random() * 50) + 95}%</div>
                <div class="stat-label">Success Rate</div>
            </div>
            <div class="stat-card">
                <div class="stat-value">${Math.floor(Math.random() * 20) + 5}</div>
                <div class="stat-label">Active Models</div>
            </div>
        </div>
    `;
    
    // Add footer
    const footerHtml = `
        <div class="footer">
            <p>© 2023 CampaignForge AI • Enterprise Marketing Solution</p>
            <p>Powered by RAG Technology and Advanced Language Models</p>
        </div>
    `;
    
    // Insert the header at the top of the main container
    const mainContainer = document.querySelector('.gradio-container > div:first-child');
    if (mainContainer) {
        mainContainer.insertAdjacentHTML('afterbegin', headerHtml);
        mainContainer.insertAdjacentHTML('beforeend', footerHtml);
        
        // Find the first main-container and insert stats after it
        const firstMainContainer = document.querySelector('.main-container');
        if (firstMainContainer) {
            firstMainContainer.insertAdjacentHTML('afterend', statsHtml);
        }
    }
    
    // Add processing animation to generate button
    const generateBtn = document.querySelector('button.primary');
    if (generateBtn) {
        generateBtn.addEventListener('click', function() {
            const outputBox = document.querySelector('[data-testid="textbox"]');
            if (outputBox) {
                outputBox.innerHTML = '<div class="processing-bar"><div class="progress" id="progress-bar"></div></div><p class="pulse">Generating your campaign... This may take a moment</p>';
                
                let width = 0;
                const interval = setInterval(function() {
                    if (width >= 90) {
                        clearInterval(interval);
                    } else {
                        width += 5;
                        const progressBar = document.getElementById('progress-bar');
                        if (progressBar) {
                            progressBar.style.width = width + '%';
                        }
                    }
                }, 200);
            }
        });
    }
    
    // Add tab navigation
    const tabsHtml = `
        <div class="tabs">
            <div class="tab active" data-tab="generate">Generate Campaign</div>
            <div class="tab" data-tab="history">Campaign History</div>
            <div class="tab" data-tab="settings">Settings</div>
        </div>
    `;
    
    const inputArea = document.querySelector('.main-container .gr-box');
    if (inputArea) {
        inputArea.insertAdjacentHTML('beforebegin', tabsHtml);
        
        // Add tab click handlers
        const tabs = document.querySelectorAll('.tab');
        tabs.forEach(tab => {
            tab.addEventListener('click', function() {
                tabs.forEach(t => t.classList.remove('active'));
                this.classList.add('active');
                
                // For demo purposes, just show an alert
                alert('Switching to ' + this.textContent + ' tab. This would load different content in a real application.');
            });
        });
    }
    
    // Add dark mode toggle
    const darkModeToggle = document.createElement('button');
    darkModeToggle.textContent = '🌙 Dark Mode';
    darkModeToggle.classList.add('secondary');
    darkModeToggle.style.marginLeft = 'auto';
    darkModeToggle.style.marginRight = '20px';
    
    darkModeToggle.addEventListener('click', function() {
        document.body.classList.toggle('dark');
        this.textContent = document.body.classList.contains('dark') ? '☀️ Light Mode' : '🌙 Dark Mode';
    });
    
    const header = document.querySelector('.header');
    if (header) {
        header.appendChild(darkModeToggle);
    }
    
    // Add model indicator to output
    const outputLabel = document.querySelector('[aria-label="Generated Campaign Script"]');
    if (outputLabel) {
        const modelName = document.querySelector('.gr-dropdown__select')?.value || 'openai/gpt-oss-120b';
        const modelBadge = document.createElement('span');
        modelBadge.classList.add('model-badge');
        modelBadge.textContent = modelName;
        outputLabel.appendChild(modelBadge);
    }
    
    // Update model badge when model changes
    const modelDropdown = document.querySelector('.gr-dropdown__select');
    if (modelDropdown) {
        modelDropdown.addEventListener('change', function() {
            const modelBadge = document.querySelector('.model-badge');
            if (modelBadge) {
                modelBadge.textContent = this.value;
            }
        });
    }
});
</script>
"""

# -------------------------------
# Gradio UI
# -------------------------------
with gr.Blocks(fill_height=True, theme=gr.themes.Soft(), css=custom_css) as demo:
    # Add custom JavaScript
    gr.HTML(js)
    
    # Top badges
    with gr.Row(elem_classes="badge-container"):
        gr.HTML("""
            <div style="display: flex; justify-content: center; align-items: center; gap: 15px; flex-wrap: wrap;">
                <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>
                <a href="#" target="_blank">
                    <img src="https://img.shields.io/static/v1?label=Enterprise&message=Grade&color=%23000080&labelColor=%23ffa500&logo=azurepipelines&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", elem_classes="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"
                    )
                    
                gr.Markdown("---")
                gr.Markdown("### 📈 Recent Activity")
                gr.HTML("""
                    <div style="font-size: 14px; line-height: 1.6;">
                    <p>✅ Campaign for Tech Startup - 2 hours ago</p>
                    <p>✅ Holiday Promotion - 5 hours ago</p>
                    <p>✅ Product Launch - Yesterday</p>
                    </div>
                """)

        # Main chat area
        with gr.Column(scale=3):
            with gr.Group(elem_classes="main-container"):
                gr.Markdown("## 💬 RAG-powered Creative Campaign Assistant")
                gr.Markdown("Generate compelling marketing campaigns using our AI-powered platform with retrieval-augmented generation technology.")

                query = gr.Textbox(label="Enter campaign idea / theme", lines=2, placeholder="e.g., A summer promotion for eco-friendly water bottles targeting millennials")
                output = gr.Textbox(label="Generated Campaign Script", lines=10)

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

                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")
                    
            with gr.Group(elem_classes="main-container"):
                gr.Markdown("## 📊 Campaign Performance")
                gr.HTML("""
                    <div style="display: flex; justify-content: space-between; margin: 15px 0;">
                        <div style="text-align: center;">
                            <div style="font-size: 24px; font-weight: bold; color: #2563eb;">87%</div>
                            <div style="font-size: 14px;">Engagement Rate</div>
                        </div>
                        <div style="text-align: center;">
                            <div style="font-size: 24px; font-weight: bold; color: #10b981;">2.4x</div>
                            <div style="font-size: 14px;">ROI Improvement</div>
                        </div>
                        <div style="text-align: center;">
                            <div style="font-size: 24px; font-weight: bold; color: #f59e0b;">14%</div>
                            <div style="font-size: 14px;">Conversion Rate</div>
                        </div>
                    </div>
                """)

    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()