#!/usr/bin/env python3
"""Loom harness — the search half of Loom Spark 3.

The model never searches. It decides a lookup is needed and writes the query:

    <lookup>france</lookup>

This script does the rest: searches Wikipedia, finds the ONE sentence most likely to
hold the answer, hands it back as a <result>, and lets the model answer from it.

    python3 harness.py "what's the capital of france"
    python3 harness.py                          # interactive
    python3 harness.py --no-tools "who are you"
    python3 harness.py --show "who wrote hamlet"   # print what was searched and read

How it finds the answer, and why each step exists (all measured on live questions):
  * searches the model's query AND the subject it can see in your question —
    "whats the capital of france" searched as-is returns "Capital city" and "Das Kapital"
  * prefers the real article over lists, films, albums and disambiguation pages
  * reads the article's intro first, and further only when the intro has no answer of
    the right kind (a height with a unit, a year, a number, a name)
  * strips brackets and pronunciation guides, so real text looks like training text
  * hands back ONE sentence. A 340-character window found the answer more often but the
    model misread it four times in five; one sentence doubled the final score (15% -> 30%)

Swap search() for anything you like — the contract is text in, one sentence out.
Wikipedia needs no API key. Stdlib only.
"""
from __future__ import annotations

import argparse, json, re, ssl, sys, time, urllib.error, urllib.parse, urllib.request

try:                                    # macOS system Python often lacks a CA bundle
    import certifi
    SSL_CTX = ssl.create_default_context(cafile=certifi.where())
except Exception:
    SSL_CTX = ssl.create_default_context()

OLLAMA = "http://localhost:11434/api/generate"
MODEL = "hf.co/textilelabs/Loom-Tapestry-3"
API = "https://en.wikipedia.org/w/api.php?"
# Wikipedia returns 403 without a descriptive User-Agent.
UA = {"User-Agent": "LoomHarness/3.0 (Textile Labs; https://huggingface.co/textilelabs)"}
LOOKUP = re.compile(r"<lookup>(.*?)</lookup>", re.S)
_cache: dict = {}

# ------------------------------------------------------------------- the model
def loom(prompt: str, n: int = 64) -> str:
    body = json.dumps({"model": MODEL, "prompt": prompt, "raw": True, "stream": False,
                       "options": {"temperature": 0, "num_predict": n,
                                   "stop": ["<|eot|>", "<user>", "<result>"]}}).encode()
    req = urllib.request.Request(OLLAMA, data=body, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=120) as r:
        return json.load(r)["response"].strip()

# ------------------------------------------------------------------ wikipedia
def _get(params: dict) -> dict:
    key = json.dumps(params, sort_keys=True)
    if key in _cache:
        return _cache[key]
    for attempt in range(3):
        try:
            with urllib.request.urlopen(urllib.request.Request(
                    API + urllib.parse.urlencode(params), headers=UA),
                    context=SSL_CTX, timeout=20) as r:
                _cache[key] = json.load(r)
                return _cache[key]
        except urllib.error.HTTPError as e:
            if e.code == 429:
                time.sleep(3 * (attempt + 1)); continue
            raise
    raise RuntimeError("Wikipedia rate limit")

def search(q: str, n: int = 3) -> list:
    return [h["title"] for h in _get({"action": "query", "list": "search", "srsearch": q,
                                      "format": "json", "srlimit": n})["query"]["search"]]

def _extract(title: str, intro: bool) -> str:
    p = {"action": "query", "prop": "extracts", "explaintext": 1, "titles": title,
         "format": "json", "redirects": 1}
    if intro:
        p["exintro"] = 1
    return next(iter(_get(p)["query"]["pages"].values())).get("extract", "") or ""

# ------------------------------------------------------------------ the finder
SENT = re.compile(r"(?<=[.!?])\s+(?=[A-Z0-9])")
PAREN = re.compile(r"\s*\([^()]*\)")
HEADING = re.compile(r"^\s*=+[^=]+=+\s*$", re.M)
STOP = set(("what whats who whos whom whose when where which why how is are was were be the a an "
            "of in on to for does did do by from with as at and or that this it its there tell me "
            "please can you many much").split())
ATTR = set(("capital city height tall high elevation population largest biggest smallest longest "
            "shortest tallest highest deepest first last symbol chemical language languages spoken "
            "legs year date end ended sink sank invented inventor discovered discovery developed "
            "wrote written author painted painter president founded born died age old size area "
            "distance speed").split())
JUNK = re.compile(r"^(lists? of|outline of|index of|timeline of)\b|\((film|album|song|band|"
                  r"novel|play|tv series|musical|opera|video game|book|composition|poem)\)|"
                  r"\bdisambiguation\b", re.I)

def keywords(t: str) -> list:
    return [w for w in re.findall(r"[^\W_]+", t.lower()) if w not in STOP]

def subject(question: str) -> str:
    kw = keywords(question)
    return " ".join(k for k in kw if k not in ATTR) or " ".join(kw)

def clean(t: str) -> str:
    prev = None
    while prev != t:
        prev, t = t, PAREN.sub("", t)
    return re.sub(r"\s+", " ", t.replace(" ,", ",")).strip()

def _hard(q: str, s: str) -> float:
    """The answer is of the right KIND: a height with a unit, a year, a number, a name."""
    b = 0.0
    if re.search(r"\b(how tall|how high|height|elevation)\b", q):
        b += 2.0 if re.search(r"\d[\d,.]*\s*(m|metres|meters|ft|feet|km)\b", s) else 0
    if re.search(r"\b(when|what year|which year|what date)\b", q):
        b += 2.0 if re.search(r"\b(1\d{3}|20\d{2})\b", s) else 0
    if re.search(r"\b(how many|how much|population|number of)\b", q):
        b += 1.5 if re.search(r"\d", s) else 0
    if re.search(r"\bwho\b", q):
        b += 1.5 if re.search(r"\b[A-Z][a-z]+ [A-Z][a-z]+", s) else 0
    if re.search(r"\bsymbol\b", q):
        b += 2.0 if re.search(r"\bsymbol\b", s, re.I) else 0
    if re.search(r"\bcapital\b", q):
        b += 2.0 if re.search(r"\bcapital\b", s, re.I) else 0
    return b

def _kind(q: str, s: str) -> float:
    b, sl = _hard(q, s), s.lower()
    if re.search(r"\b(how tall|how high|height|elevation)\b", q):
        b += 1.5 if re.search(r"\b(summit|elevation|height|above sea level|highest|stands)\b", sl) else -0.5
    if re.search(r"\b(end|ended|finish|finished)\b", q):
        b += 1.5 if re.search(r"\b(ended|end of|surrender|surrendered|concluded|finished)\b", sl) else -0.5
    if re.search(r"\bpopulation\b", q):
        b += 2.0 if re.search(r"\d{1,3}(,\d{3})+|\d+(\.\d+)?\s*(million|billion)", s) else -1.0
    if re.search(r"\b(invent|invented|inventor|discovered|wrote|painted|composed|founded)\b", q):
        b += 1.0 if re.search(r"\b[A-Z][a-z]+ (?:[A-Z][a-z]+ )?[A-Z][a-z]+\b", s) else 0.0
    return b

def find(query: str, question: str) -> tuple:
    """One sentence most likely to hold the answer, and the article it came from."""
    q = question.lower()
    subj = subject(question)
    pool = {}
    for tq in dict.fromkeys(x for x in (query.strip(), subj, " ".join(keywords(question))) if x):
        for rank, t in enumerate(search(tq, 3)):
            tl = t.lower()
            s = (4.0 if tl in (subj, query.strip().lower()) else 2.0 if subj and tl.startswith(subj) else 0.0)
            s += -4.0 if JUNK.search(t) else 0.0
            pool[t] = max(pool.get(t, -1e9), s - 0.3 * rank)
    qk = list(dict.fromkeys(keywords(question) + keywords(query)))
    top = sorted(pool.items(), key=lambda x: -x[1])[:3]
    typed = bool(re.search(r"\b(how tall|how high|height|elevation|when|what year|which year|"
                           r"how many|how much|population|who|symbol|capital)\b", q))
    best = (-1e9, "", "")
    for intro in (True, False):
        found_kind = False
        for title, ps in top:
            raw = _extract(title, intro)
            if re.search(r"\b(may|can) refer to\b", raw[:400]):
                continue
            body = clean(HEADING.sub(" ", raw))
            sents = [s.strip() for s in SENT.split(body) if 20 < len(s.strip()) < 600]
            for i, s in enumerate(sents[: 14 if intro else 90]):
                sc = ps + sum(1.0 for k in qk if k in s.lower()) + _kind(q, s) + (0.5 if i < 3 else 0.0)
                if sc > best[0]:
                    best = (sc, s, title)
                    found_kind = _hard(q, s) > 0
        if best[1] and (not typed or found_kind):
            break                       # the intro held an answer of the right kind
    return best[1], best[2]

# ------------------------------------------------------------------ the loop
def ask(message: str, tools: bool = True, show: bool = False) -> str:
    convo = f"<tools:{'on' if tools else 'off'}>\n<user>\n{message.strip()}\n<|eot|>\n<loom>\n"
    first = loom(convo)
    m = LOOKUP.search(first)
    if not m:
        return first
    query = m.group(1).strip()
    try:
        result, source = find(query, message)
    except Exception as e:
        # Never feed an error in as if it were a result — the model will answer from it.
        return f"[harness] lookup failed for {query!r}: {e}"
    if not result:
        return f"[harness] nothing found for {query!r}"
    if show:
        print(f"    [searched: {query!r}]\n    [read from {source}: {result[:150]}]")
    return loom(convo + first + f"<|eot|>\n<result>\n{result}\n<|eot|>\n<loom>\n", n=48)

def main() -> int:
    global MODEL
    ap = argparse.ArgumentParser(description="Loom Spark 3 harness")
    ap.add_argument("message", nargs="*")
    ap.add_argument("--no-tools", action="store_true", help="chat only, no lookups")
    ap.add_argument("--show", action="store_true", help="print the query and the sentence read")
    ap.add_argument("--model", default=MODEL)
    a = ap.parse_args()
    MODEL = a.model
    if a.message:
        print(ask(" ".join(a.message), not a.no_tools, a.show)); return 0
    print(f"Loom harness — {MODEL} (tools {'off' if a.no_tools else 'on'}, ctrl-c to quit)\n")
    while True:
        try:
            msg = input("you  > ").strip()
        except (EOFError, KeyboardInterrupt):
            print(); return 0
        if msg:
            print(f"loom > {ask(msg, not a.no_tools, a.show)}\n")

if __name__ == "__main__":
    sys.exit(main())