import spaces
import gradio as gr
import time
from pathlib import Path

from huggingface_hub import hf_hub_download


INDEX_PATH = Path("data/index/bge_m3_dense.faiss")


def ensure_faiss_index():
    """Ensure the runtime has the actual FAISS binary, not a pointer file."""

    print("[BIS Compass] Checking FAISS index...")

    if INDEX_PATH.exists():
        with INDEX_PATH.open("rb") as f:
            header = f.read(32)

        print(
            f"[BIS Compass] Local FAISS size: "
            f"{INDEX_PATH.stat().st_size / 1024 / 1024:.2f} MB"
        )
        print(f"[BIS Compass] Local FAISS header: {header[:16]!r}")

        # A Git-LFS pointer starts with:
        # version https://git-lfs.github.com/spec/v1
        if not header.startswith(b"version "):
            print("[BIS Compass] Local FAISS file looks valid.")
            return

        print("[BIS Compass] Detected Git-LFS pointer. Downloading real index...")
    else:
        print("[BIS Compass] FAISS index is missing. Downloading...")

    downloaded = hf_hub_download(
        repo_id="SpaceShark/bis-compass-backend",
        filename="data/index/bge_m3_dense.faiss",
        repo_type="space",
        force_download=True,
    )

    INDEX_PATH.parent.mkdir(parents=True, exist_ok=True)

    # Copy the actual binary into the location expected by Retriever.
    import shutil

    shutil.copy2(downloaded, INDEX_PATH)

    print(
        f"[BIS Compass] Downloaded FAISS index: "
        f"{INDEX_PATH.stat().st_size / 1024 / 1024:.2f} MB"
    )

    with INDEX_PATH.open("rb") as f:
        print(f"[BIS Compass] New FAISS header: {f.read(16)!r}")


ensure_faiss_index()

from src.offline_guard import enforce_offline_if_cached

enforce_offline_if_cached()

# IMPORTANT:
# ZeroGPU expects GPU-dependent models to be created at module scope.
# The spaces.GPU decorator will provide the real GPU when inference runs.
from src.retrieval.retriever import Retriever

print("[BIS Compass] Loading retriever...")
t0 = time.perf_counter()

retriever = Retriever()

print(f"[BIS Compass] Retriever ready in {time.perf_counter() - t0:.1f}s")


@spaces.GPU
def search(query: str):
    """Run BIS Compass retrieval on the ZeroGPU."""

    query = (query or "").strip()

    if not query:
        return "Please enter a search query."

    t0 = time.perf_counter()

    try:
        hits = retriever.search(query)
    except Exception as e:
        print(f"[BIS Compass] Search error: {e}")
        raise

    latency = time.perf_counter() - t0

    if not hits:
        return "No matching BIS standards found."

    lines = [
        f"### Search results for: `{query}`",
        "",
        f"**Latency:** `{latency * 1000:.0f} ms`",
        "",
    ]

    for hit in hits:
        lines.append(f"### {hit.rank}. {hit.is_code}")
        lines.append(f"**{hit.title}**")
        lines.append("")

        if hit.scope:
            lines.append(f"> {hit.scope}")
            lines.append("")

        lines.append(
            f"**Rerank score:** `{hit.rerank_score:.4f}`  \n"
            f"**RRF score:** `{hit.rrf_score:.4f}`"
        )

        if hit.categories:
            lines.append(f"  \n**Categories:** {', '.join(hit.categories)}")

        lines.append("")
        lines.append("---")
        lines.append("")

    return "\n".join(lines)


# ------------------------------------------------------------------
# Gradio UI
# ------------------------------------------------------------------

with gr.Blocks(
    title="BIS Compass",
    theme=gr.themes.Base(),
) as demo:
    gr.Markdown(
        """
# 🧭 BIS Compass

### BIS Standard Retrieval

Describe the product, material, application, or requirement you are
looking for, and BIS Compass will retrieve the most relevant Indian
Standards.

**Pipeline:** BM25 + bge-m3 → RRF → bge-reranker-v2-m3
"""
    )

    query = gr.Textbox(
        label="Search query",
        placeholder="e.g. cement for concrete construction",
        lines=2,
    )

    search_button = gr.Button(
        "🔍 Search",
        variant="primary",
    )

    output = gr.Markdown()

    search_button.click(
        fn=search,
        inputs=query,
        outputs=output,
        concurrency_limit=1,
    )

    query.submit(
        fn=search,
        inputs=query,
        outputs=output,
        concurrency_limit=1,
    )

    gr.Markdown(
        """
---
**BIS Compass** · 559 standards · Hybrid retrieval · ZeroGPU
"""
    )


if __name__ == "__main__":
    demo.launch()
# rebuild Sun Aug  9 12:10:10 AM IST 2026