"""Retrieval over the slides of this lecture and its appendix: BM25, optionally embeddings and their fusion.

Each slide (its text) is one chunk.

    python3 search.py "why is the middle of a long context used worse"
    python3 search.py "previous_response_id" --show-prompt
    python3 search.py "do chatbots overlook facts buried halfway through many pages" \
        --embed text-embedding-embeddinggemma-300m \
        --query-prefix "task: search result | query: " --doc-prefix "title: none | text: "

--embed calls an OpenAI-compatible /embeddings endpoint (default: the LM Studio
server on localhost:1234) and prints the dense ranking and the reciprocal-rank
fusion of both rankings. --deck takes local files or URLs of the slides; by
default it reads index.html and appendix.html two folders up, whichever exist.
Only the Python standard library is used.
"""

import argparse
import hashlib
import json
import math
import os
import re
import tempfile
import urllib.request
from collections import Counter
from html.parser import HTMLParser

HERE = os.path.dirname(os.path.abspath(__file__))
DECKS = [os.path.join(HERE, "..", "..", name) for name in ("index.html", "appendix.html")]


class Slides(HTMLParser):
    """Collects id, title and text of every <section> of the deck."""

    def __init__(self):
        super().__init__()
        self.slides, self.current, self.in_title, self.skip = [], None, False, 0

    def handle_starttag(self, tag, attrs):
        if tag == "section":
            self.current = {"id": dict(attrs).get("id", ""), "title": "", "text": []}
            self.slides.append(self.current)
        elif tag == "h2" and self.current is not None:
            self.in_title = True
        elif tag in ("script", "style", "svg"):
            self.skip += 1

    def handle_endtag(self, tag):
        if tag == "h2":
            self.in_title = False
        elif tag in ("script", "style", "svg"):
            self.skip -= 1

    def handle_data(self, data):
        if self.current is None or self.skip:
            return
        if self.in_title:
            self.current["title"] += data
        self.current["text"].append(data)


def load_slides(path):
    parser = Slides()
    if path.startswith(("http://", "https://")):
        with urllib.request.urlopen(path, timeout=60) as r:
            parser.feed(r.read().decode("utf-8"))
    else:
        with open(path, encoding="utf-8") as f:
            parser.feed(f.read())
    where = "appendix slide" if os.path.basename(path) == "appendix.html" else "slide"
    for number, s in enumerate(parser.slides, 1):
        s["number"] = number
        s["where"] = "%s %d" % (where, number)
        s["title"] = " ".join(s["title"].split())
        s["text"] = " ".join(" ".join(s["text"]).split())
    return parser.slides


def tokens(text):
    return re.findall(r"[a-z0-9]+(?:[-.'][a-z0-9]+)*", text.lower())


def bm25(query, docs, k1=1.2, b=0.75):
    """BM25 scores of every document for the query (the Lucene variant of the IDF)."""
    tf = [Counter(tokens(d)) for d in docs]
    lengths = [sum(c.values()) for c in tf]
    avg = sum(lengths) / len(docs)
    df = Counter(t for c in tf for t in c)
    n = len(docs)
    scores = []
    for c, length in zip(tf, lengths):
        s = 0.0
        for t in set(tokens(query)):
            if t in c:
                idf = math.log(1 + (n - df[t] + 0.5) / (df[t] + 0.5))
                s += idf * c[t] * (k1 + 1) / (c[t] + k1 * (1 - b + b * length / avg))
        scores.append(s)
    return scores


def embed(texts, model, url):
    out = []
    for i in range(0, len(texts), 32):
        body = json.dumps({"model": model, "input": texts[i:i + 32]}).encode()
        req = urllib.request.Request(url.rstrip("/") + "/embeddings", body,
                                     {"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=300) as r:
            data = json.load(r)["data"]
        out += [d["embedding"] for d in sorted(data, key=lambda d: d["index"])]
    return out


def cosine(u, v):
    return sum(a * b for a, b in zip(u, v)) / math.sqrt(sum(a * a for a in u) * sum(b * b for b in v))


def ranking(scores):
    return sorted(range(len(scores)), key=lambda i: -scores[i])


def show(name, order, scores, slides, k):
    print(f"\n{name}")
    for i in [i for i in order if scores[i] > 0][:k]:
        s = slides[i]
        print(f"  {scores[i]:8.3f}  {s['where']:>18}  {s['id']}: {s['title']}")


def main():
    p = argparse.ArgumentParser()
    p.add_argument("query")
    p.add_argument("-k", type=int, default=3, help="number of slides to show")
    p.add_argument("--embed", metavar="MODEL", help="embedding model on the server")
    p.add_argument("--url", default="http://localhost:1234/v1", help="OpenAI-compatible base URL")
    p.add_argument("--query-prefix", default="", help='e.g. "search_query: " for nomic-embed-text')
    p.add_argument("--doc-prefix", default="", help='e.g. "search_document: " for nomic-embed-text')
    p.add_argument("--show-prompt", action="store_true", help="print the prompt built from the top k slides")
    p.add_argument("--deck", nargs="+", default=[d for d in DECKS if os.path.exists(d)], help="decks to index")
    a = p.parse_args()
    if not a.deck:
        p.error("no slides found; pass --deck with a file or URL of the slides")

    # The demonstration slide quotes the queries; leave it out.
    slides = [s for d in a.deck for s in load_slides(d) if s["id"] != "demo-retrieval"]
    docs = [s["text"] for s in slides]
    scores = bm25(a.query, docs)
    order = ranking(scores)
    show("BM25 (exact words)", order, scores, slides, a.k)

    if a.embed:
        # Embedding models read at most a few thousand tokens; keep the start of each slide.
        texts = [a.doc_prefix + d[:2000] for d in docs]
        key = hashlib.sha256(json.dumps([a.embed, a.url, texts]).encode()).hexdigest()[:16]
        cache = os.path.join(tempfile.gettempdir(), f"slide-embeddings-{key}.json")
        if os.path.exists(cache):
            with open(cache) as f:
                vectors = json.load(f)
        else:
            vectors = embed(texts, a.embed, a.url)
            with open(cache, "w") as f:
                json.dump(vectors, f)
        q = embed([a.query_prefix + a.query], a.embed, a.url)[0]
        dense = [cosine(q, v) for v in vectors]
        dense_order = ranking(dense)
        show(f"Embeddings ({a.embed}, cosine similarity)", dense_order, dense, slides, a.k)
        # Reciprocal-rank fusion: sum over the two rankings of 1 / (60 + rank).
        scores = [0.0] * len(slides)
        for o in (order, dense_order):
            for rank, i in enumerate(o, 1):
                scores[i] += 1 / (60 + rank)
        order = ranking(scores)
        show("Hybrid (reciprocal-rank fusion)", order, scores, slides, a.k)

    if a.show_prompt:
        context = "\n\n".join(f"[{slides[i]['where']}] {slides[i]['text']}" for i in [i for i in order if scores[i] > 0][:a.k])
        prompt = f"Answer using only these sources, and cite the slide.\n\n{context}\n\nQuestion: {a.query}"
        print(f"\nPrompt ({len(prompt):,} characters):\n\n{prompt}")


if __name__ == "__main__":
    main()