import streamlit as st
import pickle
import time
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# ============================================================================
# Model Logic (Strictly Unchanged)
# ============================================================================
class CodePaulLite:
    def __init__(self, dataset):
        self.dataset = dataset
        self.keys = list(dataset.keys())
        self.vectorizer = TfidfVectorizer(stop_words="english")
        self.matrix = self.vectorizer.fit_transform(self.keys)

    def search_registries(self, query, top_k=10):
        if not query or not query.strip(): return []
        q_vec = self.vectorizer.transform([query.strip()])
        sims = cosine_similarity(q_vec, self.matrix).flatten()
        top_idx = sims.argsort()[::-1][:top_k]
        return [{"key": self.keys[i], "value": self.dataset[self.keys[i]], "score": float(sims[i])} for i in top_idx]

@st.cache_resource(show_spinner=False)
def load_model():
    with open("codepaul_lite_100k.pkl", "rb") as f:
        return pickle.load(f)

def set_query(q):
    st.session_state.query_widget = q

# ============================================================================
# AI Copilot Assets (Logo SVG)
# ============================================================================
ai_logo_svg = """
<svg width="80" height="80" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg">
    <defs>
        <linearGradient id="ai-grad" x1="2" y1="2" x2="22" y2="22" gradientUnits="userSpaceOnUse">
            <stop stop-color="#6366f1" />
            <stop offset="1" stop-color="#a855f7" />
        </linearGradient>
        <filter id="glow">
            <feGaussianBlur stdDeviation="1.5" result="coloredBlur"/>
            <feMerge><feMergeNode in="coloredBlur"/><feMergeNode in="SourceGraphic"/></feMerge>
        </filter>
    </defs>
    <path d="M12 2L14.85 8.65L22 12L14.85 15.35L12 22L9.15 15.35L2 12L9.15 8.65L12 2Z" fill="url(#ai-grad)" filter="url(#glow)"/>
    <circle cx="12" cy="12" r="3" fill="white" opacity="0.8"/>
</svg>
"""

# ============================================================================
# UI Styling - AI Copilot Edition
# ============================================================================
st.set_page_config(page_title="CodePaul AI Copilot", page_icon="✨", layout="wide")

st.markdown(f"""
    <style>
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;600&family=Fira+Code:wght@400;500&display=swap');

    html, body, .stApp {{
        background-color: #03040b !important;
        color: #f8fafc !important;
        font-family: 'Inter', sans-serif !important;
    }}

    /* Fond Copilote avec halo diffus */
    .stApp {{
        background: radial-gradient(circle at 50% 0%, rgba(99, 102, 241, 0.15) 0%, transparent 50%),
                    #03040b !important;
    }}

    /* Input AI Style */
    div[data-testid="stTextInput"] input {{
        background: rgba(255, 255, 255, 0.05) !important;
        border: 1px solid rgba(255, 255, 255, 0.1) !important;
        border-radius: 16px !important;
        color: #fff !important;
        padding: 1.2rem !important;
        font-size: 1.1rem !important;
        box-shadow: inset 0 0 20px rgba(0,0,0,0.5);
    }}

    /* Suggestion Chips - Carrées & Minimalistes */
    div[data-testid="stHorizontalBlock"] button {{
        background: rgba(255, 255, 255, 0.03) !important;
        border: 1px solid rgba(255, 255, 255, 0.1) !important;
        border-radius: 8px !important;
        height: 48px !important;
        color: #94a3b8 !important;
        font-size: 0.85rem !important;
        transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
    }}
    div[data-testid="stHorizontalBlock"] button:hover {{
        background: rgba(99, 102, 241, 0.1) !important;
        border-color: #6366f1 !important;
        color: white !important;
        transform: scale(1.02);
    }}

    /* Result Card - AI Copilot Style */
    .copilot-card {{
        background: rgba(255, 255, 255, 0.02);
        border: 1px solid rgba(255, 255, 255, 0.05);
        border-radius: 12px;
        margin-top: 1.5rem;
    }}

    .card-header {{
        padding: 15px 20px;
        background: linear-gradient(90deg, rgba(99, 102, 241, 0.1), transparent);
        border-bottom: 1px solid rgba(255, 255, 255, 0.05);
        display: flex;
        justify-content: space-between;
        align-items: center;
        border-radius: 12px 12px 0 0;
    }}

    .key-label {{
        font-family: 'Fira Code', monospace;
        color: #a5b4fc;
        font-size: 0.8rem;
    }}

    .ai-badge {{
        background: linear-gradient(90deg, #6366f1, #a855f7);
        color: white;
        padding: 4px 10px;
        border-radius: 6px;
        font-size: 0.65rem;
        font-weight: 700;
        text-transform: uppercase;
        letter-spacing: 1px;
        box-shadow: 0 0 15px rgba(99, 102, 241, 0.4);
    }}

    /* Custom Scrollbar for code blocks */
    .stCodeBlock {{
        background-color: rgba(0,0,0,0.3) !important;
        border: 1px solid rgba(255, 255, 255, 0.05) !important;
        border-radius: 0 0 12px 12px !important;
        border-top: none !important;
    }}
    </style>
""", unsafe_allow_html=True)

# ============================================================================
# Main UI
# ============================================================================
st.markdown(f"<div style='text-align:center; padding-top:20px;'>{ai_logo_svg}</div>", unsafe_allow_html=True)
st.markdown("<h1 style='text-align:center; font-weight:800; letter-spacing:-1px; margin-bottom:5px;'>CodePaul Copilot</h1>", unsafe_allow_html=True)
st.markdown("<p style='text-align:center; color:#64748b; margin-bottom:40px;'>Semantic code intelligence at your service.</p>", unsafe_allow_html=True)

try:
    model = load_model()
except:
    st.error("Model Error.")
    st.stop()

# Search
c1, c2, c3 = st.columns([1, 8, 1])
with c2:
    query = st.text_input("AI Search", placeholder="Ask me for any algorithm...", label_visibility="collapsed", key="query_widget")
    
    if not query:
        st.markdown("<p style='font-size:0.8rem; color:#475569; margin-top:20px;'>Quick assistance:</p>", unsafe_allow_html=True)
        chips = ["Graph Theory", "Optimization", "Sorting", "Neural Networks", "Data Structures"]
        cols = st.columns(len(chips))
        for i, chip in enumerate(chips):
            with cols[i]:
                st.button(chip, key=f"chip_{i}", on_click=set_query, args=(chip,), use_container_width=True)

# Results
if query and query.strip():
    results = model.search_registries(query.strip(), top_k=st.sidebar.slider("Results", 1, 20, 5))
    
    if results:
        for rank, r in enumerate(results, 1):
            # Problème du <div> corrigé ici : on ferme proprement chaque bloc
            is_best = (rank == 1)
            display_key = r['key'][:120] + "..." if len(r['key']) > 120 else r['key']
            badge_html = "<span class='ai-badge'>✦ Top Result</span>" if is_best else ""
            
            st.markdown(f"""
                <div class="copilot-card">
                    <div class="card-header">
                        <span class="key-label">{display_key}</span>
                        {badge_html}
                    </div>
                </div>
            """, unsafe_allow_html=True)
            
            st.code(r["value"], language="python")
            
            if len(r["key"]) > 120:
                with st.expander("Explore Logic Details"):
                    st.text(r["key"])
    else:
        st.info("I couldn't find a direct match. Try rephrasing your request.")

# Sidebar minimaliste
with st.sidebar:
    st.markdown("### Settings")
    st.caption("AI Search Parameters")
    st.divider()
    st.info("This copilot uses vector embeddings to find the most relevant code node.")