"""YouRunAI: persistent local semantic search over .txt and .md files."""
import hashlib
import json
from pathlib import Path
import numpy as np
import gradio as gr
from sentence_transformers import SentenceTransformer

docs = Path("documents")
docs.mkdir(exist_ok=True)
for name, text in {
    "returns.md": "Cedar Shop accepts unused items within 30 days. Contact support with the order number. Refunds go to the original payment method.",
    "shipping.md": "Cedar Shop ships weekdays. Standard shipping takes three to five business days. Tracking arrives by email.",
    "hours.md": "Cedar Shop support is open Monday to Friday, 9 am to 5 pm. Weekend messages are answered on Monday."
}.items():
    if not (docs / name).exists():
        (docs / name).write_text(text, encoding="utf-8")
model = SentenceTransformer("Qwen/Qwen3-Embedding-0.6B", device="cpu")
model.max_seq_length = 1024
store = Path("search-index")
store.mkdir(exist_ok=True)
passages, vectors = [], None

def fingerprint():
    return hashlib.sha256(b"".join(p.name.encode() + p.read_bytes() for p in sorted(docs.iterdir()) if p.suffix.lower() in [".txt", ".md"])).hexdigest()

def index():
    global passages, vectors
    passages = []
    for path in sorted(docs.iterdir()):
        if path.suffix.lower() not in [".txt", ".md"]:
            continue
        words = path.read_text(encoding="utf-8").split()
        for start in range(0, len(words), 160):
            passages.append({"file": path.name, "text": " ".join(words[start:start + 200])})
    if not passages:
        raise gr.Error("Add a UTF-8 .txt or .md file to documents first.")
    vectors = model.encode([p["text"] for p in passages], batch_size=4, normalize_embeddings=True)
    np.save(store / "vectors.npy", vectors)
    (store / "passages.json").write_text(json.dumps({"fingerprint": fingerprint(), "passages": passages}, ensure_ascii=False), encoding="utf-8")
    return f"Indexed {len(passages)} passages. Add or edit files, then click Rebuild index."

def search(query):
    if not query.strip():
        raise gr.Error("Enter a search question.")
    if vectors is None:
        raise gr.Error("Click Rebuild index first.")
    if json.loads((store / "passages.json").read_text(encoding="utf-8"))["fingerprint"] != fingerprint():
        raise gr.Error("Documents changed. Click Rebuild index before searching.")
    q = model.encode([query], prompt_name="query", normalize_embeddings=True)
    scores = (q @ vectors.T)[0]
    ranked = [{**passages[int(i)], "score": float(scores[i])} for i in np.argsort(-scores)[:5]]
    export = store / "candidates.json"
    export.write_text(json.dumps({"query": query, "candidates": ranked}, indent=2, ensure_ascii=False), encoding="utf-8")
    return [[n + 1, r["file"], round(r["score"], 4), r["text"]] for n, r in enumerate(ranked)], str(export)

if (store / "passages.json").exists() and (store / "vectors.npy").exists():
    saved = json.loads((store / "passages.json").read_text(encoding="utf-8"))
    if saved["fingerprint"] == fingerprint():
        passages, vectors = saved["passages"], np.load(store / "vectors.npy", allow_pickle=False)
if vectors is None:
    index()
with gr.Blocks(title="Private document search") as app:
    gr.Markdown("# Search your documents locally\nAdd UTF-8 .txt/.md files to **documents/**. This returns source passages, not AI-written answers.")
    rebuild = gr.Button("Rebuild index")
    status = gr.Textbox(value=f"Ready: {len(passages)} passages", label="Index status")
    query = gr.Textbox(value="How do I return an item?", label="Search question")
    button = gr.Button("Search", variant="primary")
    table = gr.Dataframe(headers=["Rank", "File", "Similarity", "Source passage"], interactive=False, wrap=True)
    download = gr.File(label="Download candidates for reranking")
    rebuild.click(index, outputs=status)
    button.click(search, inputs=query, outputs=[table, download])
app.launch(server_name="127.0.0.1", server_port=7860, share=False)