#!/usr/bin/env python3
"""Loom harness — runs the searches for Loom Spark 2.

The model never searches. It emits `<lookup>query</lookup>` and stops. This script
is the other half of the contract: it runs the lookup, feeds a `<result>` block
back, and lets the model answer from it.

    user question
        -> Loom (tools on)  ->  <lookup>who wrote Dracula</lookup>
        -> harness runs Wikipedia
        -> <result>...</result>
        -> Loom              ->  Bram Stoker.

Wikipedia is used because it is free and needs no API key. Swap `search()` for
anything you like — the contract is just "text in, text out".

Usage:
    python3 harness.py "who wrote Dracula"
    python3 harness.py                      # interactive
    python3 harness.py --no-tools "who are you"
"""
from __future__ import annotations

import argparse
import json
import re
import sys
import ssl
import urllib.parse
import urllib.request

# macOS system Python often ships without a usable CA bundle, so Wikipedia's TLS
# fails with CERTIFICATE_VERIFY_FAILED. Use certifi's bundle when it's available.
try:
    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-2"
LOOKUP = re.compile(r"<lookup>(.*?)</lookup>", re.S)
# Wikipedia returns 403 to requests without a descriptive User-Agent — their API
# policy requires one that identifies the client.
UA = {"User-Agent": "LoomHarness/1.0 (Textile Labs; loom harness demo)"}


def loom(prompt: str, n: int = 64) -> str:
    """One raw generation. raw=True so our exact prompt format reaches the model."""
    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()


def search(query: str, sentences: int = 3) -> str:
    """Wikipedia lookup. Returns a short passage, or '' if nothing is found."""
    api = "https://en.wikipedia.org/w/api.php?" + urllib.parse.urlencode({
        "action": "query", "format": "json", "list": "search",
        "srsearch": query, "srlimit": 1})
    try:
        with urllib.request.urlopen(urllib.request.Request(api, headers=UA),
                                    timeout=20, context=SSL_CTX) as r:
            hits = json.load(r)["query"]["search"]
        if not hits:
            return ""
        title = hits[0]["title"]
        summary = ("https://en.wikipedia.org/api/rest_v1/page/summary/"
                   + urllib.parse.quote(title, safe=""))
        with urllib.request.urlopen(urllib.request.Request(summary, headers=UA),
                                    timeout=20, context=SSL_CTX) as r:
            extract = json.load(r).get("extract", "")
    except Exception as e:
        return f"(search failed: {e})"
    parts = re.split(r"(?<=[.!?])\s+", extract)
    return " ".join(parts[:sentences]).strip()


def ask(message: str, tools: bool = True, verbose: bool = True) -> str:
    mode = "on" if tools else "off"
    convo = f"<tools:{mode}>\n<user>\n{message.strip()}\n<|eot|>\n<loom>\n"
    first = loom(convo)

    m = LOOKUP.search(first)
    if not m:
        return first                      # answered directly, no tool wanted

    query = m.group(1).strip()
    if verbose:
        print(f"    [loom wants: {query!r}]")
    result = search(query)
    if not result or result.startswith("(search failed"):
        # Never feed an error string in as if it were a result — the model will try
        # to answer from it. Fail loudly instead.
        return f"[harness] lookup failed for {query!r}: {result or 'no results'}"
    if verbose:
        print(f"    [result: {result[:100]}...]")

    convo += f"{first}<|eot|>\n<result>\n{result}\n<|eot|>\n<loom>\n"
    return loom(convo, n=48)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("message", nargs="*")
    ap.add_argument("--no-tools", action="store_true", help="chat only, no lookups")
    ap.add_argument("--quiet", action="store_true")
    ap.add_argument("--model", default=MODEL)
    args = ap.parse_args()
    globals()["MODEL"] = args.model

    if args.message:
        print(ask(" ".join(args.message), not args.no_tools, not args.quiet))
        return
    print(f"Loom harness — {MODEL}   (tools {'off' if args.no_tools else 'on'}, "
          f"ctrl-c to quit)\n")
    while True:
        try:
            msg = input("you  > ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            return
        if msg:
            print(f"loom > {ask(msg, not args.no_tools, not args.quiet)}\n")


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