#!/usr/bin/env python3
"""The merge: bonsai thinks, decider-4b judges, MetaCog wires them into one model.

    thinker: bonsai (bjev llama.cpp, OpenAI chat) .... http://localhost:8010
    judge:   decider-4b (/v1/systemone shim) ......... http://localhost:8008

Usage:
    .venv/bin/python merge.py --problem "What is 17*23?"
    .venv/bin/python merge.py --problems problems.jsonl --out results.jsonl
"""

from __future__ import annotations

import argparse
import json
import sys
import time

import re

def clean_answer(text: str) -> str:
    """Strip <think>...</think> reasoning blocks the thinker may leak."""
    return re.sub(r"<think>.*?</think>\s*", "", text, flags=re.S).strip()


from metacog import Config, MetaCog, OpenAICompatThinker, SystemOneJudge

import os

THINKER_URL = os.environ.get("THINKER_URL", "http://localhost:8010")
THINKER_MODEL = os.environ.get("THINKER_MODEL", "bonsai")
JUDGE_URL = os.environ.get("JUDGE_URL", "http://localhost:8008")


def build_thinker_judge():
    thinker = OpenAICompatThinker(THINKER_URL, model=THINKER_MODEL)
    # non-TypeSafe base_url -> no auth header; the local shim is open
    judge = SystemOneJudge(base_url=JUDGE_URL)
    return thinker, judge


def build(n_paths: int = 8, no_cascade: bool = False) -> MetaCog:
    thinker, judge = build_thinker_judge()
    config = Config(
        mode="best_of_n",
        strategy="noul",
        n_paths=n_paths,
        greedy_anchor=True,
        cascade_confidence=None if no_cascade else 0.95,
        max_tokens=2048,
        temperature=0.8,
    )
    return MetaCog(thinker, judge, config)


def main() -> int:
    ap = argparse.ArgumentParser(description="bonsai + decider-4b MetaCog merge")
    src = ap.add_mutually_exclusive_group(required=True)
    src.add_argument("--problem")
    src.add_argument("--problems")
    ap.add_argument("--n-paths", type=int, default=8)
    ap.add_argument("--no-cascade", action="store_true")
    ap.add_argument("--stop-confidence", type=float, default=0.80,
                    help="race loop early-stop bar (default 0.80)")
    ap.add_argument("--no-race", action="store_true",
                    help="use the old sequential MetaCog best_of_n instead of the race loop")
    ap.add_argument("--out", default="runs/merge_results.jsonl")
    args = ap.parse_args()

    if args.problem:
        problems = [args.problem]
    else:
        problems = [json.loads(l)["problem"] for l in open(args.problems)
                    if l.strip()]
    if not problems:
        print("no problems", file=sys.stderr)
        return 1

    mc = build(args.n_paths, args.no_cascade)
    use_race = not args.no_race
    thinker = judge = None
    if use_race:
        from race import race
        thinker, judge = build_thinker_judge()
    t0 = time.time()
    with open(args.out, "w") as out:
        for i, problem in enumerate(problems):
            t1 = time.time()
            if use_race:
                r = race(problem, thinker, judge, n_paths=args.n_paths,
                         stop_confidence=args.stop_confidence)
                answer, finished = clean_answer(r.answer), r.finished
                conf, toks, jcalls, dt = (r.confidence, r.thinker_tokens,
                                          r.judge_calls, r.seconds)
                extra = {"early_stopped": r.early_stopped,
                         "winner_source": r.winner_source,
                         "paths_judged": r.paths_judged}
            else:
                result = mc.run(problem)
                dt = time.time() - t1
                rnd = result.trace.rounds[-1] if result.trace.rounds else None
                answer, finished = clean_answer(result.answer), result.finished
                conf = round(rnd.verdict.confidence, 4) if rnd else 0.0
                toks, jcalls = result.trace.thinker_tokens, result.trace.judge_calls
                extra = {}
            row = {
                "i": i,
                "problem": problem,
                "answer": answer,
                "finished": finished,
                "judge_confidence": round(conf, 4),
                "thinker_tokens": toks,
                "judge_calls": jcalls,
                "seconds": round(dt, 2),
                **extra,
            }
            out.write(json.dumps(row) + "\n")
            out.flush()
            print(f"[{i+1}/{len(problems)}] conf={conf:.2f} "
                  f"{dt:.1f}s: {answer[:120]}", flush=True)
    print(f"done: {len(problems)} problems in {time.time()-t0:.0f}s -> {args.out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())