"""See what the modules do: the same request to the bare model (llama-server) and to the system (lfm3_server.py).

    llama-server -m weights/LFM3-1.7B-Q5_K_M.gguf -ngl 99 -np 4 -c 65536 -fa on --port 8080
    python lfm3_server.py --backend http://localhost:8080          # second terminal
    python compare.py                                               # built-in examples
    python compare.py --budget 8192 "your question"

With a limited thinking budget the bare model often runs out of tokens while still reasoning and returns no answer;
the system closes the reasoning at the budget and answers. For code it also runs the examples from the statement and
repairs the program; for SVG / HTML it renders the page and repairs what the browser reports.
"""
import argparse
import json
import re
import time
import urllib.request

EXAMPLES = [
    "Find the number of ordered pairs of positive integers (a, b) such that a + b = 1000 and neither a nor b has a "
    "zero digit. Put the final answer in \\boxed{}.",
    "Write a Python 3 program that reads an integer N (1 <= N <= 10^5) and then N integers, and prints the length of "
    "the longest strictly increasing contiguous subarray.\n\nSample Input 1\n\n5\n1 2 2 3 4\n\nSample Output 1\n\n3\n\n"
    "Sample Input 2\n\n3\n5 4 3\n\nSample Output 2\n\n1\n\nPut the complete program in a single ```python code block.",
]


def chat(url, model, question, budget):
    body = {"model": model, "messages": [{"role": "user", "content": question}], "max_tokens": budget,
            "temperature": 0.6, "top_p": 0.95, "top_k": 50, "seed": 0}
    req = urllib.request.Request(url + "/v1/chat/completions", data=json.dumps(body).encode(),
                                 headers={"Content-Type": "application/json"})
    t = time.time()
    r = json.loads(urllib.request.urlopen(req, timeout=3600).read())
    msg = r["choices"][0]["message"]
    text = msg.get("content") or ""
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.S)
    if "<think>" in text or (not text.strip() and msg.get("reasoning_content")):
        text = ""  # still thinking when the budget ran out
    return text.strip(), r["choices"][0].get("finish_reason"), r.get("lfm3"), time.time() - t


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("question", nargs="?")
    ap.add_argument("--backend", default="http://localhost:8080", help="bare llama-server")
    ap.add_argument("--system", default="http://localhost:8090", help="lfm3_server.py")
    ap.add_argument("--budget", type=int, default=8192, help="token budget for both")
    a = ap.parse_args()
    for q in [a.question] if a.question else EXAMPLES:
        print("=" * 100 + "\n" + q[:300] + "\n")
        text, fin, _, dt = chat(a.backend, "base", q, a.budget)
        print(f"--- bare model ({dt:.0f} s, finish: {fin}):\n" + (text[-1200:] if text else "(no answer: the budget ran out "
                                                                     "while the model was still thinking)") + "\n")
        text, fin, info, dt = chat(a.system, "lfm3-deliberate", q, a.budget)
        print(f"--- with modules ({dt:.0f} s, {info}):\n{text[-1200:]}\n")