import re
import numpy as np
import pandas as pd
import faiss
import gradio as gr
from pathlib import Path
from huggingface_hub import snapshot_download
from sentence_transformers import SentenceTransformer

MODEL_REPO_ID = "ElMETRICO/crosstalk-ai-full-artifacts"

print("Downloading artifacts from:", MODEL_REPO_ID)
artifact_dir = Path(snapshot_download(repo_id=MODEL_REPO_ID, repo_type="model"))

MODEL_DIR = artifact_dir / "model" / "e5_lexical_contrastive_finetuned"
INDEX_PATH = artifact_dir / "artifacts" / "trained_e5_source_to_meaning.index"
DF_PATH = artifact_dir / "artifacts" / "trained_e5_source_to_meaning_df.csv"

print("Loading fine-tuned E5 model...")
model = SentenceTransformer(str(MODEL_DIR))
model.max_seq_length = 128

print("Loading FAISS index...")
source_index = faiss.read_index(str(INDEX_PATH))

print("Loading dataframe...")
source_df = pd.read_csv(DF_PATH, encoding="utf-8-sig")
print("Rows loaded:", len(source_df))


def clean_text(x):
    x = "" if pd.isna(x) else str(x)
    x = re.sub(r"[\u200b-\u200d\ufeff]", "", x)
    x = re.sub(r"\s+", " ", x).strip()
    return x


def normalize_lookup_text(x):
    return clean_text(x).lower()


def remove_parentheses_text(x):
    x = clean_text(x)
    x = re.sub(r"\(.*?\)", "", x)
    x = re.sub(r"\s+", " ", x).strip().lower()
    return x


def count_source_files(x):
    x = "" if pd.isna(x) else str(x)
    return len([p for p in x.split("||") if p.strip()]) if x.strip() else 0


def add_quality_score(df):
    df = df.copy()

    for col in ["english_meaning", "bangla_meaning", "source_file"]:
        if col not in df.columns:
            df[col] = ""
        df[col] = df[col].fillna("").astype(str)

    if "duplicate_count" not in df.columns:
        df["duplicate_count"] = 1

    df["duplicate_count"] = pd.to_numeric(df["duplicate_count"], errors="coerce").fillna(1)
    df["has_english"] = df["english_meaning"].str.strip().ne("").astype(int)
    df["has_bangla"] = df["bangla_meaning"].str.strip().ne("").astype(int)
    df["source_file_count"] = df["source_file"].apply(count_source_files)

    df["quality_score"] = (
        df["has_english"] * 3.0 +
        df["has_bangla"] * 2.0 +
        np.log1p(df["duplicate_count"]) * 0.5 +
        np.log1p(df["source_file_count"]) * 0.5
    )

    return df


source_df = add_quality_score(source_df)

for col in ["language", "source_text", "english_meaning", "bangla_meaning", "part_of_speech"]:
    if col not in source_df.columns:
        source_df[col] = ""
    source_df[col] = source_df[col].apply(clean_text)

source_df["norm_source"] = source_df["source_text"].apply(normalize_lookup_text)
source_df["base_source"] = source_df["source_text"].apply(remove_parentheses_text)


def format_verified_result(df, query, method, top_k=10):
    result = df.copy()
    result = result.sort_values(
        by=["quality_score", "duplicate_count"],
        ascending=[False, False]
    )

    keep_cols = [
        "language",
        "source_text",
        "english_meaning",
        "bangla_meaning",
        "part_of_speech",
        "duplicate_count",
        "quality_score"
    ]

    for col in keep_cols:
        if col not in result.columns:
            result[col] = ""

    result = result[keep_cols].head(top_k).copy()
    result.insert(0, "query", query)
    result.insert(1, "score", 1.0)
    result.insert(2, "method", method)

    if result["language"].nunique() > 1 or result["source_text"].nunique() > 1:
        result["confidence"] = "high_but_ambiguous"
    else:
        result["confidence"] = "high"

    result["note"] = "Verified dictionary match."
    return result


def trained_semantic_fallback(query, top_k=5, search_k_per_language=20):
    languages = sorted(source_df["language"].dropna().unique().tolist())

    query_texts = [
        f"query: {lang} word: {query}"
        for lang in languages
    ]

    query_emb = model.encode(
        query_texts,
        batch_size=16,
        convert_to_numpy=True,
        normalize_embeddings=True,
        show_progress_bar=False
    ).astype("float32")

    scores, indices = source_index.search(query_emb, search_k_per_language)

    best_by_index = {}

    for lang_i, lang in enumerate(languages):
        for score, idx in zip(scores[lang_i], indices[lang_i]):
            idx = int(idx)
            score = float(score)

            if idx < 0:
                continue

            if idx not in best_by_index or score > best_by_index[idx]["score"]:
                best_by_index[idx] = {"score": score}

    ranked = sorted(
        best_by_index.items(),
        key=lambda x: x[1]["score"],
        reverse=True
    )[:top_k]

    if len(ranked) == 0:
        return pd.DataFrame()

    selected_indices = [idx for idx, _ in ranked]
    result = source_df.iloc[selected_indices].copy().reset_index(drop=True)

    result.insert(0, "query", query)
    result.insert(1, "score", [x["score"] for _, x in ranked])
    result.insert(2, "method", "fine_tuned_e5_semantic_fallback")

    def label_score(score):
        if score >= 0.88:
            return "medium_high_trained_semantic_candidate"
        elif score >= 0.80:
            return "medium_trained_semantic_candidate"
        else:
            return "low_trained_semantic_candidate"

    result["confidence"] = result["score"].apply(label_score)
    result["note"] = "Semantic candidate only; not a confirmed translation."

    keep_cols = [
        "query",
        "score",
        "method",
        "confidence",
        "language",
        "source_text",
        "english_meaning",
        "bangla_meaning",
        "part_of_speech",
        "duplicate_count",
        "quality_score",
        "note"
    ]

    for col in keep_cols:
        if col not in result.columns:
            result[col] = ""

    return result[keep_cols]


def safe_search(query):
    query = clean_text(query)

    if not query:
        return "⚠️ Please enter a word.", "No input provided.", pd.DataFrame()

    q_norm = normalize_lookup_text(query)
    q_base = remove_parentheses_text(query)

    exact = source_df[source_df["norm_source"] == q_norm].copy()

    if len(exact) > 0:
        result = format_verified_result(exact, query, "hybrid_exact_source_match", top_k=10)
        return (
            "✅ Verified dictionary match found.",
            f"High-confidence verified dictionary output for **{query}**.",
            result
        )

    base = source_df[source_df["base_source"] == q_base].copy()

    if len(base) > 0:
        result = format_verified_result(base, query, "hybrid_base_form_match", top_k=10)
        return (
            "✅ Verified base-form match found.",
            f"High-confidence base-form dictionary output for **{query}**.",
            result
        )

    semantic = trained_semantic_fallback(query, top_k=5, search_k_per_language=20)

    if len(semantic) == 0:
        return (
            "❌ No match found.",
            "No verified dictionary match or semantic candidate was found.",
            pd.DataFrame()
        )

    strong = semantic[semantic["score"] >= 0.88].copy()

    if len(strong) > 0:
        return (
            "🧠 Semantic candidates found.",
            "No verified dictionary match was found. Showing fine-tuned E5 semantic candidates only. These are suggestions, not confirmed translations.",
            strong
        )

    return (
        "⚠️ Low-confidence semantic candidates.",
        "No verified dictionary match was found. Semantic scores are below the safe verification threshold, so no translation is claimed.",
        semantic
    )





custom_css = """
.gradio-container {
    background: #212121 !important;
    color: #ECECEC !important;
    font-family: Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif !important;
}

#shell {
    max-width: 980px;
    margin: 0 auto;
    padding: 34px 18px 42px 18px;
}

.topbar {
    display: flex;
    justify-content: space-between;
    align-items: center;
    margin-bottom: 26px;
    color: #B4B4B4;
    font-size: 14px;
}

.brand {
    font-weight: 700;
    color: #ECECEC;
}

.model-badge {
    background: #2F2F2F;
    border: 1px solid #3A3A3A;
    color: #CFCFCF;
    padding: 7px 11px;
    border-radius: 999px;
    font-size: 12px;
}

.hero {
    text-align: center;
    padding: 22px 10px 20px 10px;
    margin-bottom: 18px;
}

.title {
    font-size: 38px;
    line-height: 1.12;
    font-weight: 750;
    letter-spacing: -0.035em;
    margin: 0;
    color: #F5F5F5;
}

.subtitle {
    max-width: 760px;
    margin: 14px auto 0 auto;
    color: #B4B4B4;
    font-size: 15px;
    line-height: 1.65;
}

.capabilities {
    display: grid;
    grid-template-columns: repeat(3, 1fr);
    gap: 12px;
    margin: 26px 0 22px 0;
}

.cap-card {
    background: #2A2A2A;
    border: 1px solid #3A3A3A;
    border-radius: 16px;
    padding: 14px 15px;
    color: #D7D7D7;
    font-size: 13px;
    line-height: 1.45;
}

.cap-card strong {
    display: block;
    color: #FFFFFF;
    font-size: 14px;
    margin-bottom: 4px;
}

.search-panel {
    background: #2F2F2F !important;
    border: 1px solid #454545 !important;
    border-radius: 24px !important;
    padding: 14px 16px 16px 16px !important;
    box-shadow: 0 14px 38px rgba(0,0,0,0.20);
    margin-bottom: 18px !important;
}

textarea, input {
    background: #2F2F2F !important;
    color: #FFFFFF !important;
    border: 1px solid #4A4A4A !important;
    border-radius: 18px !important;
    font-size: 16px !important;
}

textarea:focus, input:focus {
    border-color: #6B7280 !important;
    box-shadow: 0 0 0 2px rgba(255,255,255,0.08) !important;
}

button.send-button {
    background: #ECECEC !important;
    color: #111111 !important;
    border: 1px solid #ECECEC !important;
    border-radius: 14px !important;
    font-weight: 800 !important;
    min-height: 44px !important;
}

button.send-button:hover {
    background: #FFFFFF !important;
}

button.clear-button {
    background: #262626 !important;
    color: #D1D5DB !important;
    border: 1px solid #414141 !important;
    border-radius: 14px !important;
    font-weight: 700 !important;
    min-height: 44px !important;
}

.output-panel {
    background: transparent !important;
    border: none !important;
    padding: 0 !important;
    margin-top: 8px !important;
}

.answer-card {
    background: #2A2A2A;
    border: 1px solid #3D3D3D;
    border-radius: 18px;
    padding: 18px 20px;
    margin-bottom: 14px;
    color: #ECECEC;
    line-height: 1.65;
}

.answer-title {
    font-weight: 800;
    color: #FFFFFF;
    margin-bottom: 8px;
}

.answer-note {
    color: #C7C7C7;
    font-size: 14px;
}

.dataframe {
    border-radius: 16px !important;
    overflow: hidden !important;
    border: 1px solid #3D3D3D !important;
}

.footer {
    text-align: center;
    color: #9CA3AF;
    font-size: 12.5px;
    line-height: 1.6;
    margin-top: 24px;
}

#examples-block {
    margin-top: 10px;
}

#examples-block button {
    border-radius: 999px !important;
    background: #2A2A2A !important;
    border: 1px solid #3D3D3D !important;
    color: #D1D5DB !important;
}

.block {
    background: transparent !important;
}

label {
    color: #CFCFCF !important;
}

@media (max-width: 900px) {
    #shell {
        padding: 24px 14px 34px 14px;
    }

    .title {
        font-size: 30px;
    }

    .capabilities {
        grid-template-columns: 1fr;
    }
}
"""


EMPTY_RESULTS = pd.DataFrame(
    columns=[
        "query",
        "score",
        "method",
        "confidence",
        "language",
        "source_text",
        "english_meaning",
        "bangla_meaning",
        "part_of_speech",
        "note"
    ]
)


def clean_markdown_summary(x):
    x = "" if x is None else str(x)
    x = x.replace('<div class="result-note">', "")
    x = x.replace("</div>", "")
    return x.strip()


def chat_style_search(query):
    status, summary, result = safe_search(query)
    summary = clean_markdown_summary(summary)

    answer_html = f"""
    <div class="answer-card">
        <div class="answer-title">{status}</div>
        <div class="answer-note">{summary}</div>
    </div>
    """

    return answer_html, result


def clear_interface():
    return "", """
    <div class="answer-card">
        <div class="answer-title">Ready.</div>
        <div class="answer-note">Enter a source word to retrieve verified dictionary matches or semantic candidates.</div>
    </div>
    """, EMPTY_RESULTS


with gr.Blocks(
    css=custom_css,
    theme=gr.themes.Base(
        primary_hue="slate",
        secondary_hue="slate",
        neutral_hue="zinc"
    )
) as demo:

    with gr.Column(elem_id="shell"):
        gr.HTML(
            """
            <div class="topbar">
                <div class="brand">CrossTalk AI</div>
                <div class="model-badge">Full trained retrieval system</div>
            </div>

            <div class="hero">
                <h1 class="title">What ethnic word would you like to search?</h1>
                <div class="subtitle">
                    CrossTalk AI identifies and retrieves meanings from low-resource ethnic language dictionaries using verified lexical matching,
                    fine-tuned E5 semantic retrieval, FAISS vector search, and confidence-aware output control.
                </div>
            </div>

            <div class="capabilities">
                <div class="cap-card">
                    <strong>Verified lookup</strong>
                    Exact and base-form dictionary matches are returned first.
                </div>
                <div class="cap-card">
                    <strong>Semantic fallback</strong>
                    Fine-tuned E5 suggests candidates when verified matches are unavailable.
                </div>
                <div class="cap-card">
                    <strong>Safe output</strong>
                    Low-confidence semantic results are not claimed as confirmed translations.
                </div>
            </div>
            """
        )

        with gr.Column(elem_classes=["search-panel"]):
            query = gr.Textbox(
                label="Source word",
                placeholder="Message CrossTalk AI...",
                lines=1,
                show_label=False
            )

            with gr.Row():
                clear_btn = gr.Button("Clear", elem_classes=["clear-button"])
                submit_btn = gr.Button("Search", elem_classes=["send-button"])

            gr.Examples(
                examples=[
                    ["kəkhyáŋ"],
                    ["Hula"],
                    ["Aina"],
                    ["aam"],
                    ["bajaoo"],
                    ["unknown tribal word"]
                ],
                inputs=query,
                label="Try examples",
                elem_id="examples-block"
            )

        with gr.Column(elem_classes=["output-panel"]):
            answer = gr.HTML(
                """
                <div class="answer-card">
                    <div class="answer-title">Ready.</div>
                    <div class="answer-note">Enter a source word to retrieve verified dictionary matches or semantic candidates.</div>
                </div>
                """
            )

            results = gr.Dataframe(
                label="Retrieved results",
                value=EMPTY_RESULTS,
                interactive=False,
                wrap=True,
                max_height=470
            )

        gr.HTML(
            """
            <div class="footer">
                Exact/base-form matches are verified dictionary outputs. Fine-tuned semantic fallback results are candidate suggestions only, not confirmed translations.
            </div>
            """
        )

    submit_btn.click(
        fn=chat_style_search,
        inputs=query,
        outputs=[answer, results]
    )

    query.submit(
        fn=chat_style_search,
        inputs=query,
        outputs=[answer, results]
    )

    clear_btn.click(
        fn=clear_interface,
        inputs=None,
        outputs=[query, answer, results]
    )


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