import os
import gradio as gr
import torch
from sentence_transformers import SentenceTransformer, util
from pinecone import Pinecone

# ── Model & Pinecone ────────────────────────────────────────────────────────
device = "cpu"
model = SentenceTransformer("ShiniChien/SpouseBERT", device=device)

api_key = os.getenv("PINECONE_API_KEY")
if api_key is None:
    raise ValueError("❌ ERROR: PINECONE_API_KEY not found! Please add it to your Space Settings.")
pc = Pinecone(api_key=api_key)
index = pc.Index(name="latentlove", host="https://latentlove-9e75b0f.svc.aped-4627-b74a.pinecone.io")


# ── Helpers ─────────────────────────────────────────────────────────────────
def format_input(name: str, dob: str) -> str:
    name = name.strip()
    dob = dob.strip()
    if dob:
        return f"{name} {dob}"
    return name


def get_love_label(score: float) -> tuple[str, str]:
    if score >= 0.90:
        return "💍 Soulmates", "#ff1744"
    elif score >= 0.75:
        return "💖 Deep Connection", "#e91e63"
    elif score >= 0.60:
        return "🌹 Romantic Spark", "#f06292"
    elif score >= 0.45:
        return "💫 Kindred Spirits", "#ce93d8"
    elif score >= 0.30:
        return "🤝 Friendly Vibes", "#90caf9"
    else:
        return "🌱 Just Strangers", "#b0bec5"


# ── Tab 1: Love Match ────────────────────────────────────────────────────────
def analyze_love(name1: str, dob1: str, name2: str, dob2: str):
    if not name1.strip() or not name2.strip():
        return (
            gr.update(value="⚠️ Please enter both names to reveal your destiny!", visible=True),
            gr.update(visible=False),
            gr.update(visible=False),
        )

    text1 = format_input(name1, dob1)
    text2 = format_input(name2, dob2)

    emb1 = model.encode(text1, convert_to_tensor=True, device=device)
    emb2 = model.encode(text2, convert_to_tensor=True, device=device)
    cosine_score = float(util.cos_sim(emb1, emb2)[0][0])

    display_score = max(0.0, min(1.0, cosine_score))
    percentage = round(display_score * 100, 1)
    label, color = get_love_label(display_score)

    result_html = f"""
    <div style="text-align:center; padding: 20px;">
        <div style="font-size: 3em; margin-bottom: 10px;">{label.split()[0]}</div>
        <div style="font-size: 1.4em; font-weight: bold; color: {color}; margin-bottom: 16px;">
            {label}
        </div>
        <div style="background: linear-gradient(135deg, #1a0a0f, #2d0a1a);
                    border: 1px solid {color}44;
                    border-radius: 16px; padding: 20px; margin: 10px 0;">
            <div style="font-size: 1.1em; color: #f8bbd0; margin-bottom: 12px;">
                ✨ <em>{name1.strip()}</em> &amp; <em>{name2.strip()}</em>
            </div>
            <div style="background: #2a0a15; border-radius: 50px; height: 28px;
                        overflow: hidden; border: 1px solid {color}66; margin: 10px 0;">
                <div style="height: 100%; width: {percentage}%;
                            background: linear-gradient(90deg, #880e4f, {color});
                            border-radius: 50px; transition: width 1s ease;
                            display: flex; align-items: center; justify-content: flex-end;
                            padding-right: 10px;">
                    <span style="color: white; font-size: 0.85em; font-weight: bold;">
                        {percentage}%
                    </span>
                </div>
            </div>
            <div style="font-size: 2em; font-weight: bold; color: {color}; margin-top: 8px;">
                {percentage}% Love Match
            </div>
        </div>
        <div style="font-size: 0.8em; color: #9e9e9e; margin-top: 16px; font-style: italic;">
            Raw cosine similarity: {cosine_score:.6f}
        </div>
        <div style="font-size: 0.75em; color: #616161; margin-top: 8px;">
            ✦ Powered by SpouseBERT · For entertainment purposes only ✦
        </div>
    </div>
    """

    return (
        gr.update(value="", visible=False),
        gr.update(value=result_html, visible=True),
        gr.update(visible=False),
    )


def clear_all():
    return (
        "", "", "", "",
        gr.update(visible=False),
        gr.update(visible=False),
        gr.update(visible=True),
    )


# ── Tab 2: Destiny Search ────────────────────────────────────────────────────
def destiny_search(query: str, dob: str, filter_1word: bool, filter_2words: bool, top_k: int):
    if not query.strip():
        return build_search_html([], query, error="⚠️ Please enter a name to search!")

    # Combine name + dob
    query_text = query.strip()
    dob = dob.strip()
    if dob:
        query_text = f"{query_text} {dob}"

    query_vec = model.encode(query_text, convert_to_tensor=False, device=device).tolist()

    # Metadata filter
    if filter_1word and filter_2words:
        meta_filter = {"num_words": {"$in": [1, 2]}}
    elif filter_1word:
        meta_filter = {"num_words": {"$eq": 1}}
    elif filter_2words:
        meta_filter = {"num_words": {"$eq": 2}}
    else:
        meta_filter = None

    query_kwargs = dict(
        vector=query_vec,
        top_k=top_k + 5,  # small buffer to compensate filtering
        include_metadata=True,
    )

    if meta_filter:
        query_kwargs["filter"] = meta_filter

    response = index.query(**query_kwargs)
    matches = response.get("matches", [])

    # 🔥 Remove input name from results
    query_clean = query.strip().lower()
    filtered_matches = []

    for m in matches:
        text = m.get("metadata", {}).get("text", "").strip().lower()
        if text != query_clean:
            filtered_matches.append(m)

    # Trim back to top_k
    filtered_matches = filtered_matches[:top_k]

    return build_search_html(filtered_matches, query)


def build_search_html(matches: list, query: str, error: str = "") -> str:
    if error:
        return f'<div style="text-align:center; padding:30px; color:#f48fb1; font-size:1.1em;">{error}</div>'

    if not matches:
        return (
            '<div style="text-align:center; padding:30px; color:#9e9e9e; font-style:italic;">'
            "✦ No matches found in the latent space ✦</div>"
        )

    rows_html = ""
    for i, m in enumerate(matches):
        score = float(m.get("score", 0))
        text  = m.get("metadata", {}).get("text", m.get("id", "Unknown"))
        nw    = m.get("metadata", {}).get("num_words", "?")
        label, color = get_love_label(score)
        pct   = round(min(max(score, 0.0), 1.0) * 100, 1)
        rank_emoji = ["🥇", "🥈", "🥉"][i] if i < 3 else f"#{i+1}"

        rows_html += f"""
        <div style="
            display: flex; align-items: center; gap: 16px;
            background: linear-gradient(135deg, rgba(30,5,20,0.9), rgba(50,10,30,0.85));
            border: 1px solid {color}33;
            border-radius: 14px; padding: 14px 18px; margin-bottom: 10px;
            box-shadow: 0 2px 12px rgba(233,30,99,0.08);
        ">
            <div style="font-size:1.6em; min-width:36px; text-align:center;">{rank_emoji}</div>
            <div style="flex:1;">
                <div style="font-size:1.1em; font-weight:bold; color:#fce4ec; font-family:'Playfair Display',serif;">
                    {text}
                </div>
                <div style="font-size:0.8em; color:#9e9e9e; margin-top:2px;">
                    {nw}-word name &nbsp;·&nbsp; {label}
                </div>
            </div>
            <div style="text-align:right; min-width:80px;">
                <div style="font-size:1.3em; font-weight:bold; color:{color};">{pct}%</div>
                <div style="
                    background:#2a0a15; border-radius:50px; height:6px;
                    overflow:hidden; border:1px solid {color}44; margin-top:4px; width:80px;
                ">
                    <div style="
                        height:100%; width:{pct}%;
                        background:linear-gradient(90deg,#880e4f,{color});
                        border-radius:50px;
                    "></div>
                </div>
            </div>
        </div>
        """

    return f"""
    <div style="padding: 4px 0;">
        <div style="
            text-align:center; color:#f48fb1; font-style:italic;
            font-size:0.95em; margin-bottom:16px; opacity:0.8;
        ">
            ✨ Destiny matches for <strong style="color:#fce4ec;">"{query}"</strong>
            &nbsp;·&nbsp; {len(matches)} result(s)
        </div>
        {rows_html}
    </div>
    """


def clear_search():
    return "", "", gr.update(value=""), True, True, 10


# ── CSS ──────────────────────────────────────────────────────────────────────
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Playfair+Display:ital,wght@0,400;0,700;1,400&family=Lato:wght@300;400&display=swap');

body, .gradio-container {
    background: radial-gradient(ellipse at top, #1a0010 0%, #0d0008 60%, #000000 100%) !important;
    font-family: 'Lato', sans-serif !important;
    min-height: 100vh;
}

.main-title {
    text-align: center;
    font-family: 'Playfair Display', serif !important;
    font-size: 3.2em !important;
    background: linear-gradient(135deg, #ff80ab, #f48fb1, #fce4ec, #f48fb1, #ff80ab);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    margin-bottom: 4px !important;
    letter-spacing: 2px;
}

.subtitle {
    text-align: center;
    color: #f48fb1 !important;
    font-size: 1.05em !important;
    font-style: italic;
    margin-bottom: 24px !important;
    opacity: 0.85;
}

.card-box {
    background: linear-gradient(135deg, rgba(30,5,20,0.95), rgba(50,10,30,0.9)) !important;
    border: 1px solid rgba(244,143,177,0.25) !important;
    border-radius: 20px !important;
    padding: 24px !important;
    backdrop-filter: blur(12px);
    box-shadow: 0 8px 32px rgba(233,30,99,0.15), inset 0 1px 0 rgba(255,255,255,0.05) !important;
}

.section-label {
    color: #f48fb1 !important;
    font-family: 'Playfair Display', serif !important;
    font-size: 1.15em !important;
    margin-bottom: 12px !important;
    letter-spacing: 1px;
}

label span {
    color: #f8bbd0 !important;
    font-size: 0.95em !important;
}

input[type="text"], textarea {
    background: rgba(20,5,15,0.8) !important;
    border: 1px solid rgba(244,143,177,0.3) !important;
    border-radius: 12px !important;
    color: #fce4ec !important;
    font-size: 1em !important;
    padding: 10px 14px !important;
    transition: border-color 0.3s ease, box-shadow 0.3s ease !important;
}

input[type="text"]:focus, textarea:focus {
    border-color: #f48fb1 !important;
    box-shadow: 0 0 0 3px rgba(244,143,177,0.15) !important;
    outline: none !important;
}

input::placeholder {
    color: #6d3a4a !important;
    font-style: italic;
}

.analyze-btn {
    background: linear-gradient(135deg, #880e4f, #c2185b, #e91e63) !important;
    border: none !important;
    border-radius: 50px !important;
    color: white !important;
    font-size: 1.1em !important;
    font-family: 'Playfair Display', serif !important;
    letter-spacing: 1px !important;
    padding: 14px 32px !important;
    cursor: pointer !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 20px rgba(233,30,99,0.4) !important;
    width: 100% !important;
}

.analyze-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 8px 30px rgba(233,30,99,0.6) !important;
    background: linear-gradient(135deg, #ad1457, #d81b60, #f06292) !important;
}

.clear-btn {
    background: transparent !important;
    border: 1px solid rgba(244,143,177,0.3) !important;
    border-radius: 50px !important;
    color: #f48fb1 !important;
    font-size: 0.9em !important;
    padding: 10px 24px !important;
    cursor: pointer !important;
    transition: all 0.3s ease !important;
    width: 100% !important;
}

.clear-btn:hover {
    border-color: #f48fb1 !important;
    background: rgba(244,143,177,0.08) !important;
}

.result-box {
    background: linear-gradient(135deg, rgba(20,5,15,0.98), rgba(40,8,25,0.95)) !important;
    border: 1px solid rgba(244,143,177,0.3) !important;
    border-radius: 20px !important;
    min-height: 200px;
    box-shadow: 0 0 40px rgba(233,30,99,0.1) !important;
}

.hearts-divider {
    text-align: center;
    color: #c2185b;
    font-size: 1.2em;
    letter-spacing: 8px;
    margin: 8px 0;
    opacity: 0.6;
}

.footer-note {
    text-align: center;
    color: #6d3a4a !important;
    font-size: 0.78em !important;
    margin-top: 16px !important;
    font-style: italic;
}

/* Tab styling */
.tab-nav button {
    font-family: 'Playfair Display', serif !important;
    color: #f48fb1 !important;
    font-size: 1em !important;
    border-bottom: 2px solid transparent !important;
    transition: all 0.3s ease !important;
}

.tab-nav button.selected {
    color: #fce4ec !important;
    border-bottom: 2px solid #e91e63 !important;
}

/* Checkbox styling */
.gr-checkbox label span {
    color: #f8bbd0 !important;
}

/* Slider styling */
.gr-slider input[type=range] {
    accent-color: #e91e63 !important;
}
"""

# ── UI ───────────────────────────────────────────────────────────────────────
with gr.Blocks(css=custom_css, title="LatentLove 💕") as demo:

    gr.HTML('<div class="main-title">💕 LatentLove</div>')
    gr.HTML('<div class="subtitle">Mapping the Geometry of Human Connection · Powered by SpouseBERT</div>')
    gr.HTML('<div class="hearts-divider">♥ ♥ ♥</div>')

    # ════════════════════════════════════════════════════════════════════════
    with gr.Tabs():

        # ── TAB 1: Love Match ─────────────────────────────────────────────
        with gr.Tab("💕 Love Match"):

            with gr.Row():
                with gr.Column(scale=1, elem_classes="card-box"):
                    gr.HTML('<div class="section-label">🌸 Person A</div>')
                    name1 = gr.Textbox(label="Full Name", placeholder="e.g. Taylor Swift", lines=1)
                    dob1  = gr.Textbox(label="Date of Birth (optional)", placeholder="dd/mm/yyyy", lines=1)

                with gr.Column(scale=0, min_width=60):
                    gr.HTML(
                        '<div style="display:flex;align-items:center;justify-content:center;'
                        'height:100%;font-size:2em;color:#c2185b;font-family:serif;">💞</div>'
                    )

                with gr.Column(scale=1, elem_classes="card-box"):
                    gr.HTML('<div class="section-label">🌸 Person B</div>')
                    name2 = gr.Textbox(label="Full Name", placeholder="e.g. Travis Kelce", lines=1)
                    dob2  = gr.Textbox(label="Date of Birth (optional)", placeholder="dd/mm/yyyy", lines=1)

            gr.HTML('<div style="height:12px"></div>')

            with gr.Row():
                with gr.Column(scale=2):
                    analyze_btn = gr.Button("✨ Reveal Our Connection ✨", elem_classes="analyze-btn")
                with gr.Column(scale=1):
                    clear_btn = gr.Button("↺ Clear", elem_classes="clear-btn")

            gr.HTML('<div style="height:8px"></div>')

            error_msg          = gr.HTML(visible=False)
            result_html_out    = gr.HTML(visible=False, elem_classes="result-box")
            result_placeholder = gr.HTML(
                value=(
                    '<div style="text-align:center;padding:40px;color:#6d3a4a;font-style:italic;">'
                    "✦ Enter two names above and let the embeddings reveal their story ✦</div>"
                ),
                visible=True,
            )

        # ── TAB 2: Destiny Search ─────────────────────────────────────────
        with gr.Tab("🔮 Destiny Search"):

            with gr.Column(elem_classes="card-box"):
                gr.HTML('<div class="section-label">🔮 Who is destined for you?</div>')

                search_query = gr.Textbox(
                    label="Enter a Name",
                    placeholder="e.g. Leonardo DiCaprio",
                    lines=1,
                )

                search_dob = gr.Textbox(
                    label="Date of Birth (optional)",
                    placeholder="dd/mm/yyyy",
                    lines=1,
                )

                with gr.Row():
                    cb_1word  = gr.Checkbox(label="1-word names", value=True)
                    cb_2words = gr.Checkbox(label="2-word names", value=True)

                top_k_slider = gr.Slider(
                    minimum=1,
                    maximum=50,
                    value=10,
                    step=1,
                    label="Top K results",
                )

                with gr.Row():
                    with gr.Column(scale=2):
                        search_btn = gr.Button("🔮 Reveal My Destiny ✨", elem_classes="analyze-btn")
                    with gr.Column(scale=1):
                        clear_search_btn = gr.Button("↺ Clear", elem_classes="clear-btn")

            gr.HTML('<div style="height:8px"></div>')

            search_results = gr.HTML(
                value=(
                    '<div style="text-align:center;padding:40px;color:#6d3a4a;font-style:italic;">'
                    "✦ Enter a name above and discover who the latent space has written for you ✦</div>"
                ),
                elem_classes="result-box",
            )

    # ── Footer ────────────────────────────────────────────────────────────
    gr.HTML(
        '<div class="footer-note">'
        "LatentLove uses a BERT-based embedding model trained on real-world relationship data. "
        "Results are for entertainment purposes only and do not predict actual human destiny. 🌹"
        "</div>"
    )

    # ── Events ────────────────────────────────────────────────────────────
    analyze_btn.click(
        fn=analyze_love,
        inputs=[name1, dob1, name2, dob2],
        outputs=[error_msg, result_html_out, result_placeholder],
    )

    clear_btn.click(
        fn=clear_all,
        inputs=[],
        outputs=[name1, dob1, name2, dob2, error_msg, result_html_out, result_placeholder],
    )

    search_btn.click(
        fn=destiny_search,
        inputs=[search_query, search_dob, cb_1word, cb_2words, top_k_slider],
        outputs=[search_results],
    )

    clear_search_btn.click(
        fn=clear_search,
        inputs=[],
        outputs=[search_query, search_dob, search_results, cb_1word, cb_2words, top_k_slider],
    )


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