app.py

import os
import re
import time
import threading
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed

import streamlit as st
import duckdb
import pandas as pd

# ── Config ─────────────────────────────────────────────────────────────────────
HF_TOKEN   = os.environ.get("HF_TOKEN", "")
HF_DATASET = "hf://datasets/styhero/styhero_dataset/data"
BASE_DIR   = os.path.dirname(os.path.abspath(__file__))
_L1        = os.path.join(BASE_DIR, "mobile1_index.parquet")
_L2        = os.path.join(BASE_DIR, "mobile2_index.parquet")
MOBILE1_IDX = _L1 if os.path.exists(_L1) else "hf://datasets/styhero/styhero_dataset/data/mobile1_index.parquet"
MOBILE2_IDX = _L2 if os.path.exists(_L2) else "hf://datasets/styhero/styhero_dataset/data/mobile2_index.parquet"

# Allowlist for column names that go into f-strings (can't be SQL params)
_ALLOWED_COLS = {"mobile1", "mobile2"}

# Mobile number validation: digits only, 10–15 chars (covers Indian + intl. format)
_MOBILE_RE = re.compile(r"^\d{10,15}$")

# ── Connection pool (per-thread reuse) ────────────────────────────────────────
_local = threading.local()

def get_con():
    """Reuse one DuckDB connection per thread. Connection setup costs 50-150ms
    so reusing it across queries in the same request is a big win."""
    con = getattr(_local, "con", None)
    if con is not None:
        return con
    con = duckdb.connect()
    con.execute("INSTALL httpfs; LOAD httpfs;")
    con.execute("CREATE SECRET IF NOT EXISTS hf_secret (TYPE HUGGINGFACE, TOKEN ?)", [HF_TOKEN])
    con.execute("SET threads=4;")
    con.execute("SET enable_progress_bar=false;")
    con.execute("SET enable_object_cache=true;")   # cache parquet metadata across queries
    con.execute("SET http_keep_alive=true;")
    con.execute("SET http_retries=3;")
    _local.con = con
    return con

# ── Helpers ────────────────────────────────────────────────────────────────────

def fmt_ms(ms):
    return f"{ms/1000:.2f}s" if ms >= 1000 else f"{ms:.0f}ms"

def clean(val):
    if val is None or (isinstance(val, float) and pd.isna(val)):
        return ""
    s = str(val).strip()
    if s.lower() in ("nan", "none", ""):
        return ""
    if s.endswith(".0"):
        try:
            s = str(int(float(s)))
        except Exception:
            pass
    return ", ".join(p.strip() for p in s.split("!") if p.strip())

def sanitize_df(df):
    df = df.copy()
    for col in ["name", "fname", "email", "id", "address", "circle"]:
        if col in df.columns:
            df[col] = df[col].apply(clean)
    return df

# ── Queries ────────────────────────────────────────────────────────────────────

def query_both_indexes(target):
    """Combined UNION ALL query over both index files in a single round trip.
    Parameterized to prevent SQL injection."""
    con = get_con()
    t0 = time.perf_counter()
    df = con.execute("""
        SELECT mobile1, mobile2, row_group_id, row_start, row_end, filename, 'mobile1' AS src
        FROM read_parquet(?)
        WHERE mobile1 = ?
        UNION ALL
        SELECT mobile1, mobile2, row_group_id, row_start, row_end, filename, 'mobile2' AS src
        FROM read_parquet(?)
        WHERE mobile2 = ?
    """, [MOBILE1_IDX, target, MOBILE2_IDX, target]).fetchdf()
    return df, (time.perf_counter() - t0) * 1000

def _fetch_file_batch(filename, ranges, target):
    """Fetch all matching row-groups for a single parquet file in one query.
    `ranges` is a list of (row_start, row_end, search_col) tuples."""
    # Validate search columns against allowlist before string interpolation
    for _, _, col in ranges:
        if col not in _ALLOWED_COLS:
            raise ValueError(f"Invalid search column: {col!r}")

    con = get_con()
    hf_path = f"{HF_DATASET}/{filename.split('/')[-1]}"

    # Build OR conditions; row_start/row_end are ints (we cast on input), col is allowlisted
    or_parts = []
    params = [hf_path]
    for rs, re_, col in ranges:
        or_parts.append(f"(file_row_number BETWEEN {int(rs)} AND {int(re_)} AND {col} = ?)")
        params.append(target)

    sql = f"""
        SELECT mobile1, mobile2, email, id, name, fname, address, circle
        FROM read_parquet(?, file_row_number=true)
        WHERE {' OR '.join(or_parts)}
    """
    t0 = time.perf_counter()
    df = con.execute(sql, params).fetchdf()
    return sanitize_df(df), (time.perf_counter() - t0) * 1000

# ── Cached search ─────────────────────────────────────────────────────────────

@st.cache_data(ttl=300, max_entries=200, show_spinner=False)
def search(target: str):
    """Top-level cached search. Returns (combined_df, idx_ms, fetch_ms, errors).
    TTL kept short because results contain PII — don't want them sitting in
    memory indefinitely."""
    errors = []
    t_start = time.perf_counter()

    idx_df, idx_ms = query_both_indexes(target)

    if idx_df.empty:
        return pd.DataFrame(), idx_ms, 0.0, errors

    # Group row ranges by file so we make 1 HTTP query per file, not per row-group
    by_file = defaultdict(list)
    for _, rg in idx_df[["filename", "row_start", "row_end", "src"]].drop_duplicates().iterrows():
        by_file[str(rg["filename"])].append(
            (int(rg["row_start"]), int(rg["row_end"]), str(rg["src"]))
        )

    all_dfs = []
    if by_file:
        with ThreadPoolExecutor(max_workers=min(8, len(by_file))) as ex:
            futures = {
                ex.submit(_fetch_file_batch, fname, ranges, target): fname
                for fname, ranges in by_file.items()
            }
            for f in as_completed(futures):
                try:
                    df, _ = f.result()
                    if not df.empty:
                        all_dfs.append(df)
                except Exception as e:
                    errors.append(str(e))

    fetch_ms = (time.perf_counter() - t_start) * 1000

    if not all_dfs:
        return pd.DataFrame(), idx_ms, fetch_ms, errors

    combined = pd.concat(all_dfs, ignore_index=True).drop_duplicates()
    return combined, idx_ms, fetch_ms, errors

# ── SVG icons ──────────────────────────────────────────────────────────────────

def ico(path, color="#8b8bff", size=14):
    return f'<svg width="{size}" height="{size}" viewBox="0 0 24 24" fill="none" stroke="{color}" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" style="flex-shrink:0;margin-top:1px">{path}</svg>'

ICO_PHONE   = ico('<rect x="5" y="2" width="14" height="20" rx="2"/><line x1="12" y1="18" x2="12.01" y2="18"/>')
ICO_EMAIL   = ico('<rect x="2" y="4" width="20" height="16" rx="2"/><path d="m22 7-8.97 5.7a1.94 1.94 0 0 1-2.06 0L2 7"/>')
ICO_ID      = ico('<rect x="2" y="7" width="20" height="14" rx="2"/><path d="M16 21V5a2 2 0 0 0-2-2h-4a2 2 0 0 0-2 2v16"/>')
ICO_PIN     = ico('<path d="M20 10c0 6-8 12-8 12s-8-6-8-12a8 8 0 0 1 16 0Z"/><circle cx="12" cy="10" r="3"/>')
ICO_CIRCLE  = ico('<circle cx="12" cy="12" r="10"/><path d="M12 8v4l3 3"/>')
ICO_PERSON  = ico('<path d="M20 21v-2a4 4 0 0 0-4-4H8a4 4 0 0 0-4 4v2"/><circle cx="12" cy="7" r="4"/>')

# ── Result card ────────────────────────────────────────────────────────────────

def field_html(ico_svg, label, value, full=False):
    if not value:
        return ""
    span = 'style="grid-column:1/-1"' if full else ""
    return f"""
<div class="c-field" {span}>
  <div class="c-label">{ico_svg}<span>{label}</span></div>
  <div class="c-value">{value}</div>
</div>"""

def show_card(row, idx):
    name    = clean(row.get("name"))    or "Unknown"
    fname   = clean(row.get("fname"))   or ""
    mobile1 = clean(row.get("mobile1")) or ""
    mobile2 = clean(row.get("mobile2")) or ""
    email   = clean(row.get("email"))   or ""
    rid     = clean(row.get("id"))      or ""
    address = clean(row.get("address")) or ""
    circle  = clean(row.get("circle"))  or ""

    copy_text = "\n".join(filter(None, [
        f"Name: {name}",
        f"Father: {fname}"    if fname   else "",
        f"Mobile1: {mobile1}" if mobile1 else "",
        f"Mobile2: {mobile2}" if mobile2 else "",
        f"Email: {email}"     if email   else "",
        f"ID: {rid}"          if rid     else "",
        f"Address: {address}" if address else "",
        f"Circle: {circle}"   if circle  else "",
    ]))
    copy_safe = copy_text.replace("\\","\\\\").replace("`","\\`").replace("${","\\${")

    initials = (name[0] if name and name != "Unknown" else "?").upper()

    card = f"""<!DOCTYPE html><html><head><meta charset="utf-8">
<style>
*{{box-sizing:border-box;margin:0;padding:0}}
body{{background:transparent;font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",sans-serif}}
.card{{
  background:rgba(255,255,255,.04);
  backdrop-filter:blur(24px);
  -webkit-backdrop-filter:blur(24px);
  border:1px solid rgba(255,255,255,.08);
  border-radius:18px;
  padding:20px;
  position:relative;
  box-shadow:0 20px 60px rgba(0,0,0,.5),inset 0 1px 0 rgba(255,255,255,.06);
  margin-bottom:0;
}}
@media(max-width:480px){{
  .c-grid{{grid-template-columns:1fr!important}}
  .c-full{{grid-column:1!important}}
}}
.copy-btn{{
  position:absolute;top:14px;right:14px;
  background:rgba(124,111,255,.15);
  color:#a0a0ff;
  border:1px solid rgba(124,111,255,.3);
  border-radius:8px;padding:5px 14px;
  font-size:.76rem;font-weight:600;cursor:pointer;
  transition:background .15s,color .15s;
  -webkit-tap-highlight-color:transparent;
}}
.copy-btn:hover,.copy-btn:active{{background:linear-gradient(135deg,#7c6fff,#c084fc);color:#fff;border-color:transparent}}
.header{{display:flex;align-items:center;gap:12px;padding-right:90px;margin-bottom:14px}}
.avatar{{
  width:40px;height:40px;border-radius:50%;flex-shrink:0;
  background:linear-gradient(135deg,#7c6fff,#c084fc);
  display:flex;align-items:center;justify-content:center;
  font-size:1.1rem;font-weight:700;color:#fff;
}}
.c-name{{font-size:1.05rem;font-weight:700;color:#f0f0ff}}
.c-fname{{font-size:.78rem;color:rgba(200,200,255,.45);margin-top:2px}}
.c-grid{{
  display:grid;
  grid-template-columns:repeat(auto-fit,minmax(170px,1fr));
  gap:10px 18px;
}}
.c-field{{display:flex;flex-direction:column;gap:3px}}
.c-label{{
  display:flex;align-items:center;gap:4px;
  font-size:.62rem;text-transform:uppercase;letter-spacing:.09em;
  color:rgba(200,200,255,.4);
}}
.c-value{{font-size:.87rem;color:rgba(240,240,255,.82);word-break:break-word}}
</style></head><body>
<div class="card">
  <button class="copy-btn" id="cb{idx}" onclick="
    navigator.clipboard.writeText(`{copy_safe}`)
      .then(()=>{{document.getElementById('cb{idx}').textContent='✓ Copied';
                  setTimeout(()=>document.getElementById('cb{idx}').textContent='Copy',1500)}})
      .catch(()=>document.getElementById('cb{idx}').textContent='Failed')
  ">Copy</button>

  <div class="header">
    <div class="avatar">{initials}</div>
    <div>
      <div class="c-name">{name}</div>
      {"<div class='c-fname'>S/O: " + fname + "</div>" if fname else ""}
    </div>
  </div>

  <div class="c-grid">
    {field_html(ICO_PHONE, "Mobile 1", mobile1)}
    {field_html(ICO_PHONE, "Mobile 2", mobile2)}
    {field_html(ICO_EMAIL, "Email", email)}
    {field_html(ICO_ID,    "ID", rid)}
    {field_html(ICO_CIRCLE,"Circle", circle)}
    {field_html(ICO_PIN,   "Address", address, full=True)}
  </div>
</div>
<script>
const card = document.querySelector('.card');
const resize = () => {{
  const h = card.getBoundingClientRect().height;
  window.parent.document.querySelectorAll('iframe').forEach(f => {{
    if (f.contentWindow === window) f.style.height = (h + 8) + 'px';
  }});
}};
if (window.ResizeObserver) {{
  new ResizeObserver(resize).observe(card);
}}
window.addEventListener('load', resize);
window.addEventListener('resize', resize);
setTimeout(resize, 100);
setTimeout(resize, 400);
</script>
</body></html>"""

    st.components.v1.html(card, height=380, scrolling=False)

# ── Global CSS ─────────────────────────────────────────────────────────────────

st.set_page_config(page_title="Mobile Lookup", layout="centered")

st.markdown("""
<style>
.stApp, .stApp > div {
  background: linear-gradient(145deg,#060612 0%,#0e0b24 40%,#16103a 100%) !important;
  background-attachment: fixed !important;
  min-height: 100vh;
}
#MainMenu, footer, header { display: none !important; }
.block-container {
  padding: 2.5rem 1.5rem 4rem !important;
  max-width: 680px !important;
}

.stTextInput > div,
.stTextInput > div > div {
  background: transparent !important;
  border: none !important;
  box-shadow: none !important;
  padding: 0 !important;
}

.stTextInput input {
  background: rgba(255,255,255,.06) !important;
  border: 1px solid rgba(255,255,255,.12) !important;
  border-radius: 12px !important;
  color: #e8e8ff !important;
  font-size: .97rem !important;
  padding: .7rem 1.1rem !important;
  outline: none !important;
  box-shadow: none !important;
  transition: border-color .2s !important;
}
.stTextInput input:focus {
  border-color: rgba(120,100,255,.7) !important;
  box-shadow: 0 0 0 3px rgba(100,80,255,.15) !important;
  outline: none !important;
}
.stTextInput input::placeholder { color: rgba(180,180,220,.3) !important; }
.stTextInput label { display: none !important; }
[data-testid="InputInstructions"] { display: none !important; }

[data-testid="stForm"] {
  background: rgba(255,255,255,.04) !important;
  backdrop-filter: blur(24px) !important;
  -webkit-backdrop-filter: blur(24px) !important;
  border: 1px solid rgba(255,255,255,.09) !important;
  border-radius: 20px !important;
  padding: 1.4rem !important;
  box-shadow: 0 20px 60px rgba(0,0,0,.5), inset 0 1px 0 rgba(255,255,255,.07) !important;
}

.stFormSubmitButton > button {
  background: linear-gradient(135deg,#6c5ce7 0%,#a855f7 100%) !important;
  color: #fff !important;
  border: none !important;
  border-radius: 12px !important;
  font-weight: 700 !important;
  font-size: .95rem !important;
  letter-spacing: .03em !important;
  padding: .65rem 0 !important;
  width: 100% !important;
  box-shadow: 0 4px 20px rgba(108,92,231,.4) !important;
  transition: opacity .15s, transform .1s !important;
}
.stFormSubmitButton > button:hover {
  opacity: .88 !important;
  transform: translateY(-1px) !important;
  box-shadow: 0 6px 24px rgba(108,92,231,.5) !important;
}
.stFormSubmitButton > button:active { transform: translateY(0) !important; }

.stSpinner > div {
  border-color: rgba(255,255,255,.1) !important;
  border-top-color: #7c6fff !important;
}

[data-testid="stAlert"] {
  background: rgba(255,255,255,.04) !important;
  border: 1px solid rgba(255,255,255,.08) !important;
  border-radius: 12px !important;
  backdrop-filter: blur(12px) !important;
  color: #d0d0f0 !important;
}

::-webkit-scrollbar { width: 6px; }
::-webkit-scrollbar-track { background: transparent; }
::-webkit-scrollbar-thumb { background: rgba(255,255,255,.1); border-radius: 3px; }

iframe { display: block !important; }
[data-testid="stCustomComponentV1"] { margin-bottom: 0 !important; line-height: 0; }
</style>
""", unsafe_allow_html=True)

# ── Header ──────────────────────────────────────────────────────────────────────

st.markdown("""
<div style="display:flex;align-items:center;gap:14px;margin-bottom:2rem">
  <div style="
    width:48px;height:48px;border-radius:14px;flex-shrink:0;
    background:linear-gradient(135deg,#6c5ce7,#a855f7);
    display:flex;align-items:center;justify-content:center;
    box-shadow:0 8px 24px rgba(108,92,231,.4);
  ">
    <svg width="24" height="24" viewBox="0 0 24 24" fill="none" stroke="#fff"
      stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
      <rect x="5" y="2" width="14" height="20" rx="2"/>
      <line x1="12" y1="18" x2="12.01" y2="18"/>
    </svg>
  </div>
  <div>
    <div style="font-size:1.5rem;font-weight:800;letter-spacing:-.01em;
      background:linear-gradient(90deg,#a78bfa,#e879f9);
      -webkit-background-clip:text;-webkit-text-fill-color:transparent;
      background-clip:text;line-height:1.2">Mobile Lookup</div>
  </div>
</div>
""", unsafe_allow_html=True)

# ── Search form ─────────────────────────────────────────────────────────────────

with st.form("search_form"):
    number = st.text_input(
        "Mobile number",
        placeholder="Enter mobile number — e.g. 919876543210",
        label_visibility="collapsed",
    )
    submitted = st.form_submit_button("Search", use_container_width=True, type="primary")

if not submitted or not number.strip():
    st.stop()

target = number.strip()

# Input validation — reject anything that isn't 10-15 digits
if not _MOBILE_RE.match(target):
    st.error("Please enter a valid mobile number (10–15 digits, numbers only).")
    st.stop()

# ── Search with progress ────────────────────────────────────────────────────────

with st.status("Searching...", expanded=False) as status:
    combined, idx_ms, fetch_ms, errors = search(target)

    if combined.empty:
        status.update(label="No records found", state="error", expanded=False)
        for err in errors:
            st.error(err)
        st.warning(f"No records found for **{target}**")
        st.stop()

    status.update(label=f"Done — {fmt_ms(fetch_ms)}", state="complete", expanded=False)

for err in errors:
    st.error(err)

# ── Results ─────────────────────────────────────────────────────────────────────

st.markdown(f"""
<div style="display:flex;align-items:center;gap:8px;margin-bottom:.75rem">
  <div style="
    background:rgba(74,222,128,.12);border:1px solid rgba(74,222,128,.25);
    border-radius:20px;padding:4px 14px;
    color:#4ade80;font-size:.8rem;font-weight:600
  ">{len(combined)} record(s)</div>
  <div style="color:rgba(200,200,255,.35);font-size:.75rem">
    {fmt_ms(idx_ms)} index · {fmt_ms(fetch_ms)} total
  </div>
</div>
""", unsafe_allow_html=True)

for i, (_, row) in enumerate(combined.iterrows()):
    show_card(row, i)