app.py

import streamlit as st
import requests
import json
import os
from pathlib import Path

# ─────────────────────────── Page config ──────────────────────────────────────
st.set_page_config(
    page_title="Code Search",
    page_icon="⌕",
    layout="wide",
    initial_sidebar_state="expanded",
)

# ─────────────────────────── Custom CSS ───────────────────────────────────────
st.markdown("""
<style>
@import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@300;400;500;600&family=Space+Grotesk:wght@300;400;500;600;700&display=swap');

/* ── Root theme ── */
:root {
    --bg:       #0d1117;
    --bg2:      #161b22;
    --bg3:      #21262d;
    --border:   #30363d;
    --accent:   #58a6ff;
    --accent2:  #3fb950;
    --accent3:  #f78166;
    --accent4:  #d2a8ff;
    --text:     #e6edf3;
    --text2:    #8b949e;
    --text3:    #484f58;
    --mono:     'JetBrains Mono', monospace;
    --sans:     'Space Grotesk', sans-serif;
}

/* ── Global resets ── */
html, body, [class*="css"] {
    font-family: var(--sans) !important;
    background-color: var(--bg) !important;
    color: var(--text) !important;
}
.stApp { background-color: var(--bg) !important; }

/* ── Sidebar ── */
[data-testid="stSidebar"] {
    background-color: var(--bg2) !important;
    border-right: 1px solid var(--border) !important;
}
[data-testid="stSidebar"] * { color: var(--text) !important; }

/* ── Inputs ── */
.stTextInput input, .stTextArea textarea, .stSelectbox select,
[data-testid="stTextInput"] input, [data-baseweb="input"] input {
    background-color: var(--bg3) !important;
    border: 1px solid var(--border) !important;
    border-radius: 6px !important;
    color: var(--text) !important;
    font-family: var(--mono) !important;
    font-size: 13px !important;
}
.stTextInput input:focus, .stTextArea textarea:focus {
    border-color: var(--accent) !important;
    box-shadow: 0 0 0 3px rgba(88,166,255,0.1) !important;
}

/* ── Buttons ── */
.stButton > button {
    background-color: var(--accent) !important;
    color: #0d1117 !important;
    border: none !important;
    border-radius: 6px !important;
    font-family: var(--mono) !important;
    font-weight: 600 !important;
    font-size: 13px !important;
    padding: 8px 20px !important;
    transition: all 0.2s !important;
}
.stButton > button:hover {
    background-color: #79b8ff !important;
    transform: translateY(-1px) !important;
    box-shadow: 0 4px 12px rgba(88,166,255,0.3) !important;
}
.stButton > button[kind="secondary"] {
    background-color: var(--bg3) !important;
    color: var(--text) !important;
    border: 1px solid var(--border) !important;
}
.stButton > button[kind="secondary"]:hover {
    border-color: var(--accent3) !important;
    color: var(--accent3) !important;
    background-color: rgba(247,129,102,0.1) !important;
    transform: none !important;
    box-shadow: none !important;
}

/* ── Tabs ── */
.stTabs [data-baseweb="tab-list"] {
    background-color: var(--bg2) !important;
    border-bottom: 1px solid var(--border) !important;
    gap: 0 !important;
}
.stTabs [data-baseweb="tab"] {
    background-color: transparent !important;
    color: var(--text2) !important;
    font-family: var(--mono) !important;
    font-size: 13px !important;
    padding: 12px 20px !important;
    border-radius: 0 !important;
    border-bottom: 2px solid transparent !important;
}
.stTabs [aria-selected="true"] {
    background-color: transparent !important;
    color: var(--accent) !important;
    border-bottom: 2px solid var(--accent) !important;
}
.stTabs [data-baseweb="tab-panel"] {
    background-color: var(--bg) !important;
    padding-top: 24px !important;
}

/* ── Metric cards ── */
[data-testid="metric-container"] {
    background-color: var(--bg2) !important;
    border: 1px solid var(--border) !important;
    border-radius: 8px !important;
    padding: 16px !important;
}
[data-testid="metric-container"] label {
    color: var(--text2) !important;
    font-family: var(--mono) !important;
    font-size: 11px !important;
    text-transform: uppercase !important;
    letter-spacing: 0.08em !important;
}
[data-testid="metric-container"] [data-testid="stMetricValue"] {
    color: var(--accent) !important;
    font-family: var(--mono) !important;
    font-size: 28px !important;
    font-weight: 600 !important;
}

/* ── File uploader ── */
[data-testid="stFileUploader"] {
    background-color: var(--bg2) !important;
    border: 1px dashed var(--border) !important;
    border-radius: 8px !important;
}
[data-testid="stFileUploader"]:hover {
    border-color: var(--accent) !important;
}

/* ── Success / error / info boxes ── */
.stSuccess { background-color: rgba(63,185,80,0.1) !important; border-left: 3px solid var(--accent2) !important; }
.stError   { background-color: rgba(247,129,102,0.1) !important; border-left: 3px solid var(--accent3) !important; }
.stInfo    { background-color: rgba(88,166,255,0.08) !important; border-left: 3px solid var(--accent) !important; }
.stWarning { background-color: rgba(210,168,255,0.1) !important; border-left: 3px solid var(--accent4) !important; }

/* ── JSON / code blocks ── */
.stJson, pre, code {
    background-color: var(--bg2) !important;
    border: 1px solid var(--border) !important;
    border-radius: 6px !important;
    font-family: var(--mono) !important;
    font-size: 12px !important;
}

/* ── Divider ── */
hr { border-color: var(--border) !important; margin: 24px 0 !important; }

/* ── Result cards ── */
.result-card {
    background: var(--bg2);
    border: 1px solid var(--border);
    border-radius: 8px;
    padding: 16px;
    margin-bottom: 12px;
    font-family: var(--mono);
    font-size: 13px;
    transition: border-color 0.2s;
}
.result-card:hover { border-color: var(--accent); }
.result-rank {
    display: inline-block;
    background: var(--accent);
    color: #0d1117;
    border-radius: 4px;
    padding: 2px 8px;
    font-size: 11px;
    font-weight: 700;
    margin-bottom: 8px;
}
.result-score {
    float: right;
    color: var(--accent2);
    font-size: 12px;
    font-weight: 600;
}
.result-text {
    white-space: pre-wrap;
    color: var(--text);
    line-height: 1.6;
    margin-top: 8px;
    border-top: 1px solid var(--border);
    padding-top: 8px;
}

/* ── Doc table ── */
.doc-row {
    display: flex;
    align-items: center;
    background: var(--bg2);
    border: 1px solid var(--border);
    border-radius: 6px;
    padding: 10px 16px;
    margin-bottom: 8px;
    gap: 12px;
    font-family: var(--mono);
    font-size: 13px;
}
.doc-id { color: var(--accent); flex: 1; font-weight: 500; }
.doc-chunks { color: var(--text2); font-size: 11px; }

/* ── Header ── */
.page-header {
    display: flex;
    align-items: center;
    gap: 12px;
    margin-bottom: 32px;
    padding-bottom: 20px;
    border-bottom: 1px solid var(--border);
}
.page-header h1 {
    font-family: var(--mono) !important;
    font-size: 22px !important;
    font-weight: 600 !important;
    color: var(--text) !important;
    margin: 0 !important;
    letter-spacing: -0.02em;
}
.page-header .subtitle {
    color: var(--text2);
    font-size: 13px;
    font-family: var(--mono);
    margin-top: 2px;
}
.status-dot {
    width: 8px; height: 8px;
    border-radius: 50%;
    display: inline-block;
    margin-right: 6px;
}
.status-dot.online  { background: #3fb950; box-shadow: 0 0 6px #3fb950; }
.status-dot.offline { background: #f78166; box-shadow: 0 0 6px #f78166; }

/* ── Slider ── */
[data-testid="stSlider"] .st-emotion-cache-1gv3bxi { color: var(--text2) !important; }
</style>
""", unsafe_allow_html=True)


# ─────────────────────────── Helpers ──────────────────────────────────────────
def api(method: str, path: str, **kwargs) -> requests.Response:
    base = st.session_state.get("api_url", "").rstrip("/")
    timeout = kwargs.pop("timeout", 60)
    return requests.request(method, f"{base}{path}", timeout=timeout, **kwargs)


def health_check(url: str) -> bool:
    try:
        r = requests.get(url.rstrip("/") + "/health", timeout=5)
        return r.ok and r.json().get("models_loaded", False)
    except Exception:
        return False


def render_result_card(r: dict):
    st.markdown(f"""
    <div class="result-card">
        <span class="result-rank">#{r['rank']}</span>
        <span class="result-score">score {r['score']:.4f}</span>
        <div class="result-text">{r['text']}</div>
    </div>
    """, unsafe_allow_html=True)


# ─────────────────────────── Sidebar ──────────────────────────────────────────
with st.sidebar:
    st.markdown("### ⌕ Code Search")
    st.markdown("---")

    api_url = st.text_input(
        "API URL",
        value=st.session_state.get("api_url", os.getenv("API_URL", "")),
        placeholder="https://your-space.hf.space",
        help="Your Code Search API HuggingFace Space URL",
    )
    st.session_state["api_url"] = api_url

    if api_url:
        is_up = health_check(api_url)
        dot   = "online" if is_up else "offline"
        label = "API online · models loaded" if is_up else "API offline or unreachable"
        st.markdown(
            f'<span class="status-dot {dot}"></span>'
            f'<span style="font-family:var(--mono);font-size:12px;color:var(--text2)">{label}</span>',
            unsafe_allow_html=True,
        )

    st.markdown("---")
    st.markdown(
        '<span style="font-family:var(--mono);font-size:11px;color:var(--text3)">'
        "jina-embeddings-v2-base-code<br>FAISS · AST chunking · /data persistence"
        "</span>",
        unsafe_allow_html=True,
    )


# ─────────────────────────── Header ───────────────────────────────────────────
st.markdown("""
<div class="page-header">
  <div>
    <h1>⌕ Code Search</h1>
    <div class="subtitle">semantic code search · powered by jina-embeddings-v2-base-code</div>
  </div>
</div>
""", unsafe_allow_html=True)

if not api_url:
    st.info("👈  Enter your API URL in the sidebar to get started.")
    st.stop()


# ─────────────────────────── Tabs ─────────────────────────────────────────────
t_search, t_index, t_batch, t_docs, t_embed = st.tabs([
    "🔍  Search",
    "📄  Index File",
    "📦  Batch Index",
    "🗂  Documents",
    "🧮  Embed",
])


# ══════════════════════════════════════════════════════════════════════════════
# TAB 1 — SEARCH
# ══════════════════════════════════════════════════════════════════════════════
with t_search:
    st.markdown("#### Search an indexed codebase")

    # Fetch doc list for autocomplete
    doc_ids = []
    try:
        docs_resp = api("GET", "/documents", timeout=10)
        if docs_resp.ok:
            doc_ids = [d["doc_id"] for d in docs_resp.json().get("documents", [])]
    except Exception:
        pass

    col1, col2 = st.columns([2, 1])
    with col1:
        if doc_ids:
            doc_id = st.selectbox("Document / Project ID", options=doc_ids)
        else:
            doc_id = st.text_input("Document / Project ID", placeholder="my_project")
    with col2:
        top_k = st.slider("Top K results", min_value=1, max_value=20, value=5)

    query = st.text_area(
        "Query",
        placeholder="e.g. fetch user from database\nor: async function that handles authentication",
        height=80,
    )

    if st.button("Search  ⌕", use_container_width=True):
        if not query.strip():
            st.warning("Enter a query.")
        elif not doc_id:
            st.warning("Enter a document ID.")
        else:
            with st.spinner("Searching…"):
                try:
                    r = api("POST", "/search", json={
                        "doc_id": doc_id,
                        "query":  query.strip(),
                        "top_k":  top_k,
                    })
                    if r.ok:
                        data = r.json()
                        results = data.get("results", [])
                        st.markdown(f"**{len(results)} results** for `{query[:60]}`")
                        st.markdown("---")
                        for res in results:
                            render_result_card(res)
                    else:
                        st.error(f"API error {r.status_code}: {r.text}")
                except Exception as e:
                    st.error(f"Request failed: {e}")


# ══════════════════════════════════════════════════════════════════════════════
# TAB 2 — INDEX SINGLE FILE
# ══════════════════════════════════════════════════════════════════════════════
with t_index:
    st.markdown("#### Index a single source file")

    uploaded = st.file_uploader(
        "Upload file",
        type=["py","js","ts","tsx","jsx","go","rs","java","cpp","c","cs","rb","php","md","txt"],
        help="Supported: Python, JS, TS, Go, Rust, Java, C/C++, C#, Ruby, PHP, Markdown, text",
    )

    col1, col2 = st.columns(2)
    with col1:
        custom_id = st.text_input(
            "Custom doc_id (optional)",
            placeholder="Leave blank to use filename",
        )

    if uploaded:
        st.markdown(
            f'<div class="doc-row"><span class="doc-id">{uploaded.name}</span>'
            f'<span class="doc-chunks">{uploaded.size:,} bytes</span></div>',
            unsafe_allow_html=True,
        )

    if st.button("Index File  →", use_container_width=True):
        if not uploaded:
            st.warning("Upload a file first.")
        else:
            with st.spinner(f"Indexing `{uploaded.name}`…"):
                try:
                    files   = {"file": (uploaded.name, uploaded.getvalue(), "text/plain")}
                    data    = {"doc_id": custom_id.strip()}
                    r       = api("POST", "/index", files=files, data=data, timeout=120)
                    if r.ok:
                        d = r.json()
                        c1, c2, c3 = st.columns(3)
                        c1.metric("doc_id",  d["doc_id"])
                        c2.metric("Chunks",  d["chunks_indexed"])
                        c3.metric("Status",  "✓ indexed")
                        st.success(d["message"])
                    else:
                        st.error(f"API error {r.status_code}: {r.text}")
                except Exception as e:
                    st.error(f"Request failed: {e}")


# ══════════════════════════════════════════════════════════════════════════════
# TAB 3 — BATCH INDEX
# ══════════════════════════════════════════════════════════════════════════════
with t_batch:
    st.markdown("#### Batch index an entire project")
    st.markdown(
        '<span style="font-family:var(--mono);font-size:12px;color:var(--text2)">'
        "Upload multiple files — they will all be indexed under one shared doc_id.</span>",
        unsafe_allow_html=True,
    )

    batch_files = st.file_uploader(
        "Upload files",
        accept_multiple_files=True,
        type=["py","js","ts","tsx","jsx","go","rs","java","cpp","c","cs","rb","php","md","txt"],
        key="batch_uploader",
    )

    col1, col2 = st.columns([2, 1])
    with col1:
        batch_id = st.text_input("Project doc_id", placeholder="my_project", key="batch_id")
    with col2:
        replace = st.checkbox("Replace existing index", value=True)

    if batch_files:
        st.markdown(f"**{len(batch_files)} file(s) queued:**")
        for f in batch_files:
            st.markdown(
                f'<div class="doc-row">'
                f'<span class="doc-id">{f.name}</span>'
                f'<span class="doc-chunks">{f.size:,} bytes</span>'
                f'</div>',
                unsafe_allow_html=True,
            )

    if st.button("Batch Index  →", use_container_width=True, key="batch_btn"):
        if not batch_files:
            st.warning("Upload at least one file.")
        elif not batch_id.strip():
            st.warning("Enter a project doc_id.")
        else:
            payload = {
                "doc_id":  batch_id.strip(),
                "replace": replace,
                "files": [
                    {"filename": f.name, "content": f.getvalue().decode("utf-8", errors="replace")}
                    for f in batch_files
                ],
            }
            with st.spinner(f"Indexing {len(batch_files)} files into `{batch_id}`…"):
                try:
                    r = api("POST", "/index/batch", json=payload, timeout=300)
                    if r.ok:
                        d = r.json()
                        c1, c2, c3 = st.columns(3)
                        c1.metric("doc_id",         d["doc_id"])
                        c2.metric("Files indexed",  d["files_indexed"])
                        c3.metric("Chunks indexed", d["chunks_indexed"])
                        st.success("Batch index complete!")
                    else:
                        st.error(f"API error {r.status_code}: {r.text}")
                except Exception as e:
                    st.error(f"Request failed: {e}")


# ══════════════════════════════════════════════════════════════════════════════
# TAB 4 — DOCUMENTS
# ══════════════════════════════════════════════════════════════════════════════
with t_docs:
    st.markdown("#### Indexed documents")

    col_refresh, _ = st.columns([1, 4])
    with col_refresh:
        refresh = st.button("↻ Refresh", key="refresh_docs")

    try:
        r = api("GET", "/documents", timeout=10)
        if r.ok:
            docs = r.json().get("documents", [])
            if not docs:
                st.info("No documents indexed yet. Use the Index or Batch Index tabs.")
            else:
                total_chunks = sum(d["chunks"] for d in docs)
                m1, m2 = st.columns(2)
                m1.metric("Documents", len(docs))
                m2.metric("Total chunks", total_chunks)
                st.markdown("---")

                for doc in docs:
                    col_info, col_del = st.columns([5, 1])
                    with col_info:
                        st.markdown(
                            f'<div class="doc-row">'
                            f'<span class="doc-id">📁 {doc["doc_id"]}</span>'
                            f'<span class="doc-chunks">{doc["chunks"]:,} chunks</span>'
                            f'</div>',
                            unsafe_allow_html=True,
                        )
                    with col_del:
                        if st.button("Delete", key=f"del_{doc['doc_id']}", type="secondary"):
                            try:
                                dr = api("DELETE", f"/documents/{doc['doc_id']}", timeout=10)
                                if dr.ok:
                                    st.success(f"Deleted `{doc['doc_id']}`")
                                    st.rerun()
                                else:
                                    st.error(dr.text)
                            except Exception as e:
                                st.error(str(e))
        else:
            st.error(f"API error {r.status_code}: {r.text}")
    except Exception as e:
        st.error(f"Could not reach API: {e}")


# ══════════════════════════════════════════════════════════════════════════════
# TAB 5 — EMBED
# ══════════════════════════════════════════════════════════════════════════════
with t_embed:
    st.markdown("#### Embed arbitrary texts")
    st.markdown(
        '<span style="font-family:var(--mono);font-size:12px;color:var(--text2)">'
        "Returns raw 768-dim float vectors. One text per line.</span>",
        unsafe_allow_html=True,
    )

    raw_texts = st.text_area(
        "Texts (one per line)",
        placeholder="def getUserById(id):\n    return db.query(User).filter(User.id == id).first()\nfetch user from database",
        height=160,
    )

    if st.button("Embed  →", use_container_width=True):
        texts = [t.strip() for t in raw_texts.strip().splitlines() if t.strip()]
        if not texts:
            st.warning("Enter at least one text.")
        elif len(texts) > 64:
            st.warning("Maximum 64 texts per request.")
        else:
            with st.spinner(f"Embedding {len(texts)} text(s)…"):
                try:
                    r = api("POST", "/embed", json={"texts": texts}, timeout=60)
                    if r.ok:
                        d = r.json()
                        embs = d["embeddings"]
                        st.metric("Dimensions", d["dimensions"])
                        st.markdown(f"**{len(embs)} embedding(s) returned**")
                        st.markdown("---")
                        for i, (txt, vec) in enumerate(zip(texts, embs)):
                            with st.expander(f"[{i}] `{txt[:60]}{'…' if len(txt)>60 else ''}`"):
                                preview = vec[:16]
                                st.markdown(
                                    f'<code style="font-size:11px;line-height:1.8">'
                                    f'[{", ".join(f"{v:.5f}" for v in preview)}, …]<br>'
                                    f'<span style="color:var(--text3)">dim 768 · showing first 16</span>'
                                    f'</code>',
                                    unsafe_allow_html=True,
                                )
                                if st.toggle("Show full vector", key=f"full_{i}"):
                                    st.json(vec)
                    else:
                        st.error(f"API error {r.status_code}: {r.text}")
                except Exception as e:
                    st.error(f"Request failed: {e}")