"""Interactive demo for the lethe memory server (keyword and embedding ranking)."""
import os
import time
import uuid

import gradio as gr
import httpx
from fastembed import TextEmbedding

import ranking
import rerank
from lethe import Lethe

EMBED = TextEmbedding(ranking.MODEL)  # 384-dim, runs on CPU via ONNX
BASE = "http://127.0.0.1:8787"
TOKEN = "demo"

# Embedding recall fetches this many candidates; the cross-encoder reorders them and the top 5 show.
# A wider pool than the 5 shown lets the reranker lift a right answer the embedding ranked low.
RECALL_K = 5
POOL = 20

SEED = [
    "Priya is allergic to peanuts and avoids anything cooked in groundnut oil.",
    "Priya's flight to Denver leaves on Friday at 7:40 in the morning.",
    "The team agreed to ship the billing fix before the quarterly review.",
    "Marcus prefers meetings after 2 pm because he runs in the morning.",
    "The staging database password rotates every 30 days.",
    "Priya booked a vegetarian meal for the Denver flight.",
    "Marcus will present the fraud dashboard at the quarterly review.",
]

BENCH = """| Ranking (LoCoMo, evidence in top 5) | Score |
|---|---|
| Keyword only | 0.438 |
| Embedding similarity only | 0.479 |
| Equal mix of both | 0.493 |
| Mix + neighbouring turns, hand-set weights | 0.545 |
| Learned from one score per recall | 0.568 ± 0.006 |
| **Learned from "which memories helped"** | **0.618 ± 0.002** |

Source: the repository README. Feedback in that benchmark is simulated from evidence labels, and on
LongMemEval the gain is much smaller (0.729 vs 0.727 for the equal mix). This live demo uses a different, smaller
embedding model and starting weights, and does not reproduce those numbers."""


def _wait_ready() -> None:
    for _ in range(50):
        try:
            httpx.get(BASE + "/healthz", timeout=1)
            return
        except httpx.HTTPError:
            time.sleep(0.2)


def embed(text: str) -> list[float]:
    return next(iter(EMBED.embed([text]))).tolist()


def _client(ns: str, with_embeddings: bool = True) -> Lethe:
    # Memories are always stored with an embedding; recall can leave it out to rank by keywords and neighbours only.
    return Lethe(BASE, namespace=ns, token=TOKEN, embed=embed if with_embeddings else None)


_SEEDED: set[str] = set()


def session(request: gr.Request) -> str:
    """Namespace for this visitor, seeded with the sample memories on first use (works for browser and API callers)."""
    ns = "s-" + (request.session_hash or uuid.uuid4().hex)[:10]
    if ns not in _SEEDED:
        with _client(ns) as c:
            for t in SEED:
                c.remember(t)
        _SEEDED.add(ns)
    return ns


def add_memory(text, request: gr.Request):
    if not text.strip():
        return "Nothing to store.", gr.update()
    with _client(session(request)) as c:
        c.remember(text.strip()[:1000])
    return "Stored.", ""


def do_recall(query, explore, use_embeddings, request: gr.Request):
    q = query.strip()[:500]
    if not q:
        return gr.update(choices=[]), "", None
    # Explore deliberately varies the ranking so the server can learn, so rerank only the fixed
    # ranking (explore off): the cross-encoder settles the final order the embedding leaves fragile.
    rerank_on = use_embeddings and not explore
    with _client(session(request), use_embeddings) as c:
        r = c.recall(q, k=POOL if rerank_on else RECALL_K, explore=explore)
    memories = rerank.rerank(q, r.memories)[:RECALL_K] if rerank_on else r.memories
    labels = [f"{m.score:.3f}  {m.text}" for m in memories]
    state = {"id": r.id, "ids": [m.id for m in memories], "labels": labels}
    msg = "Tick the memories that actually answer the question, then send feedback."
    return gr.update(choices=labels, value=[]), msg, state


def send_feedback(state, picked, request: gr.Request):
    if not state:
        return "Run a recall first."
    useful = [i for i, lab in zip(state["ids"], state["labels"]) if lab in (picked or [])]
    reward = len(useful) / max(len(state["ids"]), 1)
    with _client(session(request)) as c:
        c.feedback(state["id"], reward, useful=useful)
    return f"Feedback sent: {len(useful)} useful of {len(state['ids'])}, reward {reward:.2f}."


with gr.Blocks(title="lethe: agent memory that learns") as demo:
    gr.Markdown(
        "# lethe: agent memory that learns what to recall\n"
        "Real Go server running in this Space. [Code](https://github.com/shaikn6/lethe), MIT. "
        "Embeddings come from a small CPU model; switch them off to see keyword-only ranking. "
        f"With embeddings on and explore off, a cross-encoder ({rerank.MODEL}) reorders the "
        "recalled candidates so the right memory settles first."
    )
    state = gr.State(None)
    with gr.Tab("Recall and teach"):
        q = gr.Textbox(label="Question", placeholder="What can Priya not eat?")
        explore = gr.Checkbox(False, label="Explore (tries slightly different rankings so it can learn, so results vary from run to run)")
        emb = gr.Checkbox(True, label=f"Use embeddings ({ranking.MODEL}). Off = keyword and neighbour ranking only")
        go = gr.Button("Recall", variant="primary")
        note = gr.Markdown()
        picks = gr.CheckboxGroup(label="Top 5 (score, memory)", choices=[])
        fb = gr.Button("Send feedback")
        out = gr.Markdown()
        go.click(do_recall, [q, explore, emb], [picks, note, state])
        fb.click(send_feedback, [state, picks], out)
    with gr.Tab("Add memory"):
        t = gr.Textbox(label="Memory text", lines=2)
        add = gr.Button("Store")
        msg = gr.Markdown()
        add.click(add_memory, [t], [msg, t])
    with gr.Tab("Measured results"):
        gr.Markdown(BENCH)

if __name__ == "__main__":
    _wait_ready()
    ranking.apply(BASE, os.environ["LETHE_ADMIN_TOKEN"])
    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))