#!/usr/bin/env python3
"""fc_tau2_bridge.py — prompt-mode tool bridge for CloudSurf-4B-FC (W5 τ² lane).

Exposes an OpenAI-compatible /v1/chat/completions (tools in, tool_calls out)
in front of an sglang /v1/completions backend serving the merged champion.
The bridge speaks to the model in EXACTLY its trained interface:

  * prompt render byte-faithful to the vendored Gemma4Handler._format_prompt
    (prompt-mode turns, tool responses inside the model turn, thinking ON);
  * tool schemas injected via the BFCL default prompt-mode system prompt
    (the template the champion was trained and evaluated against), appended
    to the caller's own system message (τ² domain policy);
  * responses: thought channel stripped at <channel|>, full native
    <|tool_call>…<tool_call|> blocks normalized to bracket form (the
    registered handler's normalization — NO stray-closer strip), bracket
    call lists parsed into OpenAI tool_calls via ast.

Custom-scaffold disclosure: this file is the "prompt-mode tool bridge"
referenced in the τ² submission methodology notes.

Usage:  fc_tau2_bridge.py [--port 8000] [--backend http://127.0.0.1:30000]
                          [--served-name cloudsurf-4b-fc] [--selftest]
Stdlib only — no pip installs.
"""
import argparse
import ast
import json
import os
import re
import threading
import time
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

# The BFCL default prompt-mode system prompt at harness pin f7cf735
# (constants/default_prompts.py). {functions} = JSON list of function docs.
BFCL_SYSTEM_PROMPT = """You are an expert in composing functions. You are given a question and a set of possible functions. Based on the question, you will need to make one or more function/tool calls to achieve the purpose.
If none of the functions can be used, point it out. If the given question lacks the parameters required by the function, also point it out.
You should only return the function calls in your response.

If you decide to invoke any of the function(s), you MUST put it in the format of [func_name1(params_name1=params_value1, params_name2=params_value2...), func_name2(params)]
You SHOULD NOT include any other text in the response.

At each turn, you should try your best to complete the tasks requested by the user within the current turn. Continue to output functions to call until you have fulfilled the user's request to the best of your ability. Once you have no more functions to call, the system will consider the current turn complete and proceed to the next turn or task.

Here is a list of functions in JSON format that you can invoke.
{functions}
"""

NATIVE_TOOL_CALL_RE = re.compile(r"<\|tool_call>\s*call:\s*(.*?)\s*<tool_call\|>", re.DOTALL)


def format_tool_response(name: str, response: str) -> str:
    # Mirrors Gemma4Handler._format_tool_response exactly.
    return f'<|tool_response>response:{name}{{value:<|"|>{response}<|"|>}}<tool_response|>'


def calls_to_bracket(tool_calls) -> str:
    """Re-render structured OpenAI tool_calls (conversation history) into the
    bracket text the model actually emitted, so history matches training."""
    parts = []
    for tc in tool_calls or []:
        fn = tc.get("function", tc)
        name = fn.get("name", "unknown")
        try:
            args = json.loads(fn.get("arguments") or "{}")
        except Exception:
            args = {}
        rendered = ", ".join(f"{k}={args[k]!r}" for k in args)
        parts.append(f"{name}({rendered})")
    return "[" + ", ".join(parts) + "]" if parts else ""


def content_text(c) -> str:
    """Normalize OpenAI message content to plain text. Chat-completions
    clients may send content as a LIST of typed parts (the OpenAI
    content-parts format), not a string; join the text parts and ignore
    binary ones (image/audio parts cannot ride a prompt-mode render anyway)."""
    if isinstance(c, str):
        return c
    if isinstance(c, list):
        return "\n".join(
            p.get("text", "") for p in c if isinstance(p, dict) and p.get("type") == "text"
        )
    return "" if c is None else str(c)


def render_prompt(messages, tools, thinking=True) -> str:
    """Byte-faithful port of Gemma4Handler._format_prompt, with the caller's
    system message extended by the BFCL function-calling instructions."""
    functions = [t.get("function", t) for t in (tools or [])]
    fc_block = BFCL_SYSTEM_PROMPT.format(functions=json.dumps(functions))
    msgs = list(messages)
    if msgs and msgs[0]["role"] == "system":
        system_message = content_text(msgs[0].get("content")).strip() + "\n\n" + fc_block
        msgs = msgs[1:]
    else:
        system_message = fc_block

    out = ""
    if system_message or thinking:
        out += "<|turn>system\n"
        if thinking:
            out += "<|think|>\n"
        out += f"{system_message.strip()}<turn|>\n"

    i = 0
    while i < len(msgs):
        m = msgs[i]
        role = m["role"]
        content = content_text(m.get("content"))
        if role == "user":
            out += f"<|turn>user\n{content.strip()}<turn|>\n"
            i += 1
        elif role == "assistant":
            text = (content or "").strip()
            if m.get("tool_calls"):
                text = calls_to_bracket(m["tool_calls"])
            out += f"<|turn>model\n{text}"
            i += 1
            while i < len(msgs) and msgs[i]["role"] == "tool":
                tm = msgs[i]
                out += format_tool_response(tm.get("name", "unknown"), content_text(tm.get("content")))
                i += 1
            out += "<turn|>\n"
        elif role == "tool":
            out += "<|turn>model\n"
            out += format_tool_response(m.get("name", "unknown"), content_text(m.get("content")))
            out += "<turn|>\n"
            i += 1
        elif role == "system":
            out += f"<|turn>system\n{content.strip()}<turn|>\n"
            i += 1
        else:
            i += 1

    out += "<|turn>model\n"
    if not thinking:
        out += "<|channel>thought\n<channel|>"
    return out


def normalize_native(text: str) -> str:
    calls = NATIVE_TOOL_CALL_RE.findall(text)
    if not calls:
        return text
    return "[" + ", ".join(c.strip() for c in calls if c.strip()) + "]"


def parse_bracket_calls(text: str):
    """Parse '[f(a=1), g(b="x")]' into OpenAI tool_calls, or None if the text
    is not a pure bracket call list (then it is plain content)."""
    t = text.strip()
    if not (t.startswith("[") and t.endswith("]")):
        return None
    try:
        tree = ast.parse(t, mode="eval")
        if not isinstance(tree.body, ast.List):
            return None
        calls = []
        for j, el in enumerate(tree.body.elts):
            if not isinstance(el, ast.Call):
                return None
            if isinstance(el.func, ast.Attribute):
                name = ast.unparse(el.func)
            elif isinstance(el.func, ast.Name):
                name = el.func.id
            else:
                return None
            args = {}
            for kw in el.keywords:
                if kw.arg is None:
                    return None
                try:
                    args[kw.arg] = ast.literal_eval(kw.value)
                except Exception:
                    args[kw.arg] = ast.unparse(kw.value)
            if el.args:  # positional args are not valid in this format
                return None
            calls.append({
                "id": f"call_{int(time.time()*1000)%100000}_{j}",
                "type": "function",
                "function": {"name": name, "arguments": json.dumps(args)},
            })
        return calls or None
    except (SyntaxError, ValueError):
        return None


def strip_thought(raw: str):
    if "<channel|>" in raw:
        reasoning, answer = raw.rsplit("<channel|>", 1)
        return reasoning.replace("<|channel>thought", "", 1), answer
    return "", raw


class Bridge(BaseHTTPRequestHandler):
    backend = "http://127.0.0.1:30000"
    served_name = "cloudsurf-4b-fc"
    backend_model = None  # discovered from sglang /v1/models

    def log_message(self, *a):  # quiet
        pass

    def _json(self, code, obj):
        body = json.dumps(obj).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        if self.path == "/v1/models":
            self._json(200, {"object": "list", "data": [{"id": self.served_name, "object": "model"}]})
        else:
            self._json(404, {"error": "not found"})

    def do_POST(self):
        if self.path != "/v1/chat/completions":
            return self._json(404, {"error": "not found"})
        try:
            n = int(self.headers.get("Content-Length", 0))
            req = json.loads(self.rfile.read(n))
            prompt = render_prompt(req.get("messages", []), req.get("tools"))
            payload = {
                "model": Bridge.backend_model or self.served_name,
                "prompt": prompt,
                "temperature": req.get("temperature", 0.0),
                "max_tokens": req.get("max_tokens", 2048),
                "skip_special_tokens": False,
            }
            r = urllib.request.Request(
                f"{self.backend}/v1/completions",
                data=json.dumps(payload).encode(),
                headers={"Content-Type": "application/json"},
            )
            with urllib.request.urlopen(r, timeout=600) as resp:
                back = json.loads(resp.read())
            raw = back["choices"][0]["text"]
            reasoning, answer = strip_thought(raw)
            answer = normalize_native(answer).strip()
            tool_calls = parse_bracket_calls(answer)
            if not tool_calls and not answer:
                # Generation ended at/inside the thought channel. Salvage a
                # bracket call list from the full raw text before giving the
                # harness an empty answer (which scores 0 silently).
                salvaged = normalize_native(raw).strip()
                tool_calls = parse_bracket_calls(salvaged)
                if not tool_calls:
                    m = re.findall(r"\[[^\[\]]*?\([^\[\]]*?\)[^\[\]]*?\]", raw, re.DOTALL)
                    if m:
                        tool_calls = parse_bracket_calls(m[-1])
            dbg = os.getenv("BRIDGE_DEBUG_LOG")
            if dbg:
                with open(dbg, "a") as fh:
                    fh.write(json.dumps({
                        "prompt_tail": prompt[-300:],
                        "raw": raw[:3000],
                        "answer": answer[:500],
                        "n_tool_calls": len(tool_calls or []),
                    }) + "\n")
            msg = {"role": "assistant"}
            if tool_calls:
                msg["content"] = None
                msg["tool_calls"] = tool_calls
                finish = "tool_calls"
            else:
                msg["content"] = answer
                finish = "stop"
            self._json(200, {
                "id": back.get("id", "bridge"),
                "object": "chat.completion",
                "created": back.get("created", int(time.time())),
                "model": self.served_name,
                "choices": [{"index": 0, "message": msg, "finish_reason": finish}],
                "usage": back.get("usage", {}),
            })
        except Exception as e:  # surface, never hang the harness
            self._json(500, {"error": {"message": f"bridge: {type(e).__name__}: {e}"}})


def selftest():
    msgs = [
        {"role": "system", "content": "POLICY"},
        {"role": "user", "content": "What is the weather in Berkeley?"},
        {"role": "assistant", "tool_calls": [{"id": "c1", "type": "function",
            "function": {"name": "get_weather", "arguments": "{\"city\": \"Berkeley\"}"}}]},
        {"role": "tool", "name": "get_weather", "content": "72F and sunny"},
        {"role": "user", "content": "And in SF?"},
    ]
    tools = [{"type": "function", "function": {"name": "get_weather", "parameters": {
        "type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}}]
    p = render_prompt(msgs, tools)
    assert p.startswith("<|turn>system\n<|think|>\nPOLICY\n\nYou are an expert"), p[:80]
    assert "<|turn>model\n[get_weather(city='Berkeley')]<|tool_response>response:get_weather{value:<|\"|>72F and sunny<|\"|>}<tool_response|><turn|>\n" in p, p
    assert p.endswith("<|turn>user\nAnd in SF?<turn|>\n<|turn>model\n"), p[-60:]
    c = parse_bracket_calls("[get_weather(city='SF'), log(msg=\"hi there\", n=2)]")
    assert c and c[0]["function"]["name"] == "get_weather"
    assert json.loads(c[1]["function"]["arguments"]) == {"msg": "hi there", "n": 2}
    assert parse_bracket_calls("I cannot answer that.") is None
    assert parse_bracket_calls("[not a call]") is None
    n = normalize_native("<|tool_call>call: get_weather(city='SF')<tool_call|>")
    assert n == "[get_weather(city='SF')]", n
    r, a = strip_thought("<|channel>thought\nthinking...<channel|>[f(a=1)]")
    assert a == "[f(a=1)]" and "thinking" in r
    print("BRIDGE SELFTEST: ALL PASS")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=8000)
    ap.add_argument("--backend", default="http://127.0.0.1:30000")
    ap.add_argument("--served-name", default="cloudsurf-4b-fc")
    ap.add_argument("--backend-model", default=None,
                    help="model id the sglang backend expects (default: discover from /v1/models)")
    ap.add_argument("--selftest", action="store_true")
    args = ap.parse_args()
    if args.selftest:
        selftest()
        raise SystemExit(0)
    Bridge.backend = args.backend
    Bridge.served_name = args.served_name
    Bridge.backend_model = args.backend_model
    if Bridge.backend_model is None:
        try:
            with urllib.request.urlopen(f"{args.backend}/v1/models", timeout=30) as r:
                Bridge.backend_model = json.load(r)["data"][0]["id"]
        except Exception:
            Bridge.backend_model = args.served_name
    print(f"bridge on :{args.port} -> {args.backend} (backend model {Bridge.backend_model})", flush=True)
    ThreadingHTTPServer(("127.0.0.1", args.port), Bridge).serve_forever()