"""Side-by-side demo: the same small model answers a math question with no tools and with a
sandboxed run_python tool. Runs llama.cpp's llama-server on the Space's free CPU."""
import os, re, json, time, tarfile, subprocess, urllib.request, threading, glob
import gradio as gr
from huggingface_hub import hf_hub_download
import sandbox
from proxy import chat_with_python

LLAMA_TAG = "b11938"
MODEL_REPO, MODEL_FILE = "LiquidAI/LFM3-1.7B-GGUF", "LFM3-1.7B-QAD-Q4_0.gguf"
PORT = 8080
URL = os.environ.get("LLAMA_URL", f"http://127.0.0.1:{PORT}")
TOOL_REPO = "https://huggingface.co/austindixson/python-sandbox-math-tool"
SYS = "Solve the math problem. Reason step by step, then put the final answer in \\boxed{}."
_ready = threading.Event(); _status = {"msg": "starting"}

def _start_server():
    if os.environ.get("LLAMA_URL"):
        _ready.set(); return
    try:
        d = "/tmp/llama"
        if not glob.glob(f"{d}/**/llama-server", recursive=True):
            _status["msg"] = "downloading llama.cpp"
            os.makedirs(d, exist_ok=True)
            tgz = f"{d}/llama.tgz"
            urllib.request.urlretrieve(f"https://github.com/ggml-org/llama.cpp/releases/download/{LLAMA_TAG}/llama-{LLAMA_TAG}-bin-ubuntu-x64.tar.gz", tgz)
            with tarfile.open(tgz) as t: t.extractall(d)
        srv = glob.glob(f"{d}/**/llama-server", recursive=True)[0]
        _status["msg"] = "downloading model"
        gguf = hf_hub_download(MODEL_REPO, MODEL_FILE)
        _status["msg"] = "loading model"
        env = dict(os.environ, LD_LIBRARY_PATH=os.path.dirname(srv))
        subprocess.Popen([srv, "-m", gguf, "-c", "8192", "-np", "1", "-t", os.environ.get("LLAMA_THREADS", "2"),
                          "--host", "127.0.0.1", "--port", str(PORT), "--jinja"], env=env,
                         stdout=open("/tmp/llama.log", "w"), stderr=subprocess.STDOUT, cwd=os.path.dirname(srv))
        for _ in range(300):
            try:
                if json.load(urllib.request.urlopen(URL + "/health", timeout=2)).get("status") == "ok":
                    _ready.set(); _status["msg"] = "ready"; return
            except Exception: pass
            time.sleep(1)
        _status["msg"] = "model server did not start: " + open("/tmp/llama.log").read()[-500:]
    except Exception as e:
        _status["msg"] = f"startup failed: {e}"
threading.Thread(target=_start_server, daemon=True).start()

EXAMPLES = [
    "What is the sum of the digits of 2^1000?",
    "What is the remainder when 7^2026 is divided by 1000?",
    "What is 987654321 × 123456789?",
    "How many ways can you make change for one dollar using pennies, nickels, dimes and quarters?",
    "How many primes are there below 100,000?",
]

def _boxed(s):
    i = s.rfind("\\boxed{")
    if i < 0: return None
    j, d = i + 7, 1
    for k in range(j, len(s)):
        d += {"{": 1, "}": -1}.get(s[k], 0)
        if d == 0: return s[j:k]
    return None

def _one(q, use_tool, max_tokens):
    body = {"model": "lfm3", "max_tokens": int(max_tokens), "temperature": 0.0, "top_k": 1,
            "chat_template_kwargs": {"enable_thinking": False},
            "messages": [{"role": "system", "content": SYS}, {"role": "user", "content": q}]}
    t = time.time(); r = chat_with_python(URL, body, use_tool=use_tool); dt = time.time() - t
    m = r["choices"][0]["message"]; txt = re.sub(r"<think>.*?</think>", "", m.get("content") or "", flags=re.S).strip()
    x = r["x_python_tool"]; fin = r["choices"][0].get("finish_reason")
    stats = (f"**Final answer:** `{_boxed(txt) or '(no boxed answer)'}`  \n"
             f"**Generated tokens:** {r['usage']['completion_tokens']:,}  ·  **Time:** {dt:.1f} s  ·  "
             f"**Model calls:** {x['llm_calls']}  ·  **Tool calls:** {len(x['calls'])}"
             + ("  \n⚠️ hit the token cap" if fin == "length" else ""))
    calls = "\n\n".join(f"**run_python #{i}**\n```python\n{c['code']}\n```\n```text\n{c['output']}\n```" for i, c in enumerate(x["calls"], 1)) or "_no tool calls_"
    return stats, txt, calls

def compare(q, max_tokens):
    if not q.strip():
        yield "Type a question first.", "", "", "", "", ""; return
    if not _ready.wait(timeout=1):
        yield f"Model is still loading ({_status['msg']}). Try again in a minute.", "", "", "", "", ""; return
    wait = "_running..._"
    yield wait, "", "", "_waiting for the no-tools run to finish..._", "", ""
    a = _one(q, False, max_tokens)
    yield a[0], a[1], "", wait, "", ""
    b = _one(q, True, max_tokens)
    yield a[0], a[1], "", b[0], b[1], b[2]

RECORDED = open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "RECORDED.md")).read()

backend = sandbox.pick_backend()
with gr.Blocks(title="Python sandbox tool: with vs without") as demo:
    gr.Markdown(f"""# Small model, with and without a Python tool
The same model (LiquidAI **LFM3-1.7B**, QAD-Q4_0 GGUF, llama.cpp {LLAMA_TAG}, greedy, thinking off) answers your
question twice: once on its own, once with a `run_python` tool it can call up to 6 times. Runs on this Space's
free 2-vCPU CPU, so expect roughly 10–60 s per side. Tool code: [{TOOL_REPO.split('co/')[1]}]({TOOL_REPO}).

**Sandbox here:** `{backend}`: {sandbox.ISOLATION.get(backend)}. On a Mac the tool uses `sandbox-exec`
(network denied by the OS, writes only in a temp dir); on Linux with working user namespaces it uses bubblewrap.""")
    with gr.Tab("Try it"):
        with gr.Row():
            q = gr.Textbox(label="Math question", value=EXAMPLES[0], lines=2, scale=4)
            mt = gr.Slider(256, 2048, value=1024, step=128, label="Token cap per answer (demo only)", scale=1)
        gr.Examples(EXAMPLES, inputs=q)
        go = gr.Button("Run both", variant="primary")
        with gr.Row():
            with gr.Column():
                gr.Markdown("### No tools")
                s1 = gr.Markdown(); a1 = gr.Markdown(); c1 = gr.Markdown(visible=False)
            with gr.Column():
                gr.Markdown("### With run_python")
                s2 = gr.Markdown(); a2 = gr.Markdown()
                with gr.Accordion("Tool calls and outputs", open=True):
                    c2 = gr.Markdown()
        go.click(compare, [q, mt], [s1, a1, c1, s2, a2, c2], concurrency_limit=1)
        gr.Markdown("The token cap here is a demo setting. Every recorded benchmark run used a 25,000-token budget per problem.")
    with gr.Tab("Recorded results"):
        gr.Markdown(RECORDED)
        with gr.Row():
            gr.Image("fig_tokens.png", label="Tokens per problem", show_label=True)
            gr.Image("fig_aime.png", label="Hard math (AIME 2025)", show_label=True)

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch()