#!/usr/bin/env python3
"""Heuristic 1.1 race loop: parallel thought paths, judged as they finish.

Replaces the sequential best_of_n + judge-everything-at-the-end loop:

1. The greedy anchor (temperature 0) is generated and judged FIRST. If the
   judge scores it >= stop_confidence and it finished, it wins immediately --
   no sampled paths are ever generated.
2. Otherwise n_paths-1 sampled paths are generated IN PARALLEL
   (ThreadPoolExecutor over the thinker's HTTP endpoint).
3. Each path is judged the moment it completes. The first FINISHED path to
   score >= stop_confidence wins on the spot; the race is cut off at that
   instant -- queued paths are cancelled and race() returns immediately.
   (Honest cutoff semantics: threads already running a blocking HTTP
   generation can't be killed, so they finish silently in the background
   and their results are discarded. Nothing waits for them.)
4. If nothing clears the bar, the best judged finished path wins. Finished
   paths always outrank truncated ones, whatever the score.

Only the thinker's .generate() and the judge's .score() from the metacog
package are used -- the MetaCog controller is not touched.

Caveat: decider-4b's confidence scale is not calibrated (values above 1.0
have been observed), so stop_confidence is a knob to measure, not a true
probability. The 0.95 default in the MetaCog repo was measured on real Jev
(226/229 correct at >= 0.95); 0.80 here is the operator's directive and is
untested on decider-4b.
"""

from __future__ import annotations

import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass


@dataclass
class RaceResult:
    answer: str
    finished: bool
    confidence: float
    thinker_tokens: int = 0
    judge_calls: int = 0
    paths_submitted: int = 0
    paths_judged: int = 0
    early_stopped: bool = False
    winner_source: str = "sample"  # "greedy" | "sample"
    seconds: float = 0.0


def _judge_one(judge, problem: str, text: str) -> float:
    """Judge's P(correct) for a single path; 0.0 if the judge call fails."""
    try:
        return float(judge.score(problem, [text]).confidence)
    except Exception:
        return 0.0


def race(
    problem: str,
    thinker,
    judge,
    n_paths: int = 8,
    stop_confidence: float = 0.80,
    max_tokens: int = 2048,
    temperature: float = 0.8,
    max_workers: int = 8,
) -> RaceResult:
    t0 = time.time()
    lock = threading.Lock()
    thinker_tokens = 0
    judge_calls = 0

    def gen(**kw):
        nonlocal thinker_tokens
        gens = thinker.generate(problem, "", max_tokens=max_tokens, **kw)
        with lock:
            for g in gens:
                thinker_tokens += g.tokens or 0
        return gens

    def judge_of(text: str) -> float:
        nonlocal judge_calls
        with lock:
            judge_calls += 1
        return _judge_one(judge, problem, text)

    # --- 1. greedy anchor first -------------------------------------------
    greedy = gen(n=1, temperature=0.0)[0]
    g_conf = judge_of(greedy.text)
    pool = [(greedy.text, greedy.finished, g_conf, "greedy")]
    if greedy.finished and g_conf >= stop_confidence:
        return RaceResult(
            answer=greedy.text, finished=True, confidence=g_conf,
            thinker_tokens=thinker_tokens, judge_calls=judge_calls,
            paths_submitted=1, paths_judged=1,
            early_stopped=True, winner_source="greedy",
            seconds=time.time() - t0,
        )

    # --- 2. sampled paths in parallel, judged as they finish ---------------
    n_sampled = max(n_paths - 1, 0)
    early = False
    winner = None
    # NOTE: not a `with` block on purpose. Exiting a ThreadPoolExecutor
    # context manager WAITS for every worker, which would defeat the early
    # stop. Instead we shut down without waiting: queued futures are
    # cancelled, race() returns at once, and any thread already inside a
    # blocking thinker.generate() finishes on its own in the background
    # (its result is discarded). The thinker's generate() is MetaCog's
    # library code and is deliberately not modified, so true GPU-side
    # abort of an in-flight generation isn't available -- this is a
    # stop-waiting cutoff, which is what cuts user-facing latency.
    ex = ThreadPoolExecutor(max_workers=max(1, min(max_workers, n_sampled or 1)))
    try:
        futs = [ex.submit(gen, n=1, temperature=temperature)
                for _ in range(n_sampled)]
        for fut in as_completed(futs):
            try:
                g = fut.result()[0]
            except Exception:
                continue  # a dead path just never gets judged
            conf = judge_of(g.text)
            pool.append((g.text, g.finished, conf, "sample"))
            if g.finished and conf >= stop_confidence:
                winner = (g.text, True, conf, "sample")
                early = True
                break
    finally:
        ex.shutdown(wait=False, cancel_futures=True)

    # --- 3. no early stop: best judged finished path wins -------------------
    if winner is None:
        finished = [p for p in pool if p[1]]
        text, fin, conf, src = max(finished or pool, key=lambda p: p[2])
        winner = (text, fin, conf, src)

    text, fin, conf, src = winner
    with lock:
        toks, calls = thinker_tokens, judge_calls
    return RaceResult(
        answer=text, finished=fin, confidence=conf,
        thinker_tokens=toks, judge_calls=calls,
        paths_submitted=1 + n_sampled, paths_judged=len(pool),
        early_stopped=early, winner_source=src,
        seconds=time.time() - t0,
    )