"""Inference loop around the served model: execute, repair, vote. No training, no reference answers.

For a sql case the loop does what a deployed system can do with its own database:
  1. ask the model for k candidate queries (k=1: one query),
  2. check each: one SELECT, transpiles to Oracle/PostgreSQL, runs on the world's first database,
  3. a candidate that fails gets the error message back and up to `repairs` more tries,
  4. among the candidates that run, the result most of them agree on wins (ties: the earliest).
Report cases (summary + chart) pass through with one call. The runner only writes outputs
({id, sample, output, latency_s}) and a trace; `evalrun.py score` scores them with the frozen verifier.

  uv run python code/infer_loop.py --cases eval/cases/dev.jsonl --endpoint http://localhost:18000/v1 --model rft1 \
      --out runs/loop/k1r2/outputs.jsonl --n 4 --k 1 --repairs 2
"""

from __future__ import annotations

import argparse
import json
import sys
import time
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "eval"))

import verify as V
import worldlib

REPAIR = ("Bu sorgu çalışmadı. Hata: {error}\n"
          "Şemaya ve kurallara uyan, düzeltilmiş TEK bir SELECT sorgusu yaz. Cevabı yalnızca bir ```sql bloğu içinde ver.")


def check(world: str, output: str) -> tuple[str | None, list | None]:
    """(error, rows): the error a deployed system would see for this answer, or the rows it returns."""
    sql = V.extract_sql(output)
    if sql is None:
        return "cevapta SELECT sorgusu yok", None
    why = V.check_statement(sql)
    if why:
        return why, None
    try:
        return None, worldlib.run(world, worldlib.SEEDS[0], sql)[2]
    except worldlib.QueryError as e:
        return str(e), None


def result_key(rows: list) -> str:
    """Two queries vote together when they return the same cells, whatever the row and column order."""
    def cell(v):
        v = worldlib.norm(v)
        return f"{v:.9g}" if isinstance(v, float) else repr(v)

    return json.dumps(sorted(sorted(cell(v) for v in r) for r in rows))


def solve_sql(chat, messages: list, world: str, k: int = 1, repairs: int = 0, temperature: float = 0.0) -> tuple[str, dict]:
    """chat(messages, temperature) -> text. Returns the chosen answer and a trace of every step."""
    trace: dict = {"candidates": []}
    valid = []
    for i in range(k):
        out = chat(messages, temperature)
        err, rows = check(world, out)
        steps = [{"output": out, "error": err}]
        tries = 0
        while err and tries < repairs:
            tries += 1
            out = chat([*messages, {"role": "assistant", "content": out}, {"role": "user", "content": REPAIR.format(error=err[:300])}], temperature)
            err, rows = check(world, out)
            steps.append({"output": out, "error": err})
        trace["candidates"].append(steps)
        if err is None and rows is not None:
            valid.append((i, out, result_key(rows)))
    if not valid:
        trace["chosen"] = None
        return trace["candidates"][0][-1]["output"], trace
    votes: dict = {}
    for _, _, key in valid:
        votes[key] = votes.get(key, 0) + 1
    i, out, key = max(valid, key=lambda c: (votes[c[2]], -c[0]))
    trace["chosen"], trace["votes"], trace["valid"] = i, votes[key], len(valid)
    return out, trace


def make_chat(endpoint: str, model: str, max_tokens: int, top_p: float, timeout: float = 180.0):
    url = endpoint.rstrip("/") + "/chat/completions"

    def chat(messages: list, temperature: float) -> str:
        body = {"model": model, "messages": messages, "temperature": temperature, "top_p": top_p, "max_tokens": max_tokens,
                "chat_template_kwargs": {"enable_thinking": False}}
        req = urllib.request.Request(url, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"})
        for attempt in range(3):
            try:
                with urllib.request.urlopen(req, timeout=timeout) as r:
                    return json.loads(r.read())["choices"][0]["message"].get("content") or ""
            except Exception:
                if attempt == 2:
                    raise
                time.sleep(2 * (attempt + 1))
        return ""

    return chat


def run_case(chat, case: dict, sample: int, k: int, repairs: int, temperature: float) -> tuple[dict, dict]:
    t0 = time.time()
    if case.get("kind") == "sql":
        out, trace = solve_sql(chat, case["messages"], case["world"], k, repairs, temperature)
    else:
        out, trace = chat(case["messages"], temperature), {}
    row = {"id": case["id"], "sample": sample, "output": out, "latency_s": round(time.time() - t0, 3)}
    return row, {"id": case["id"], "sample": sample, **trace}


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--cases", required=True)
    ap.add_argument("--endpoint", required=True)
    ap.add_argument("--model", required=True)
    ap.add_argument("--out", required=True, help="outputs.jsonl; the trace goes next to it as trace.jsonl")
    ap.add_argument("--n", type=int, default=1, help="scored samples per case")
    ap.add_argument("--k", type=int, default=1, help="candidates per sample (voting when > 1)")
    ap.add_argument("--repairs", type=int, default=0)
    ap.add_argument("--temperature", type=float, default=1.0)
    ap.add_argument("--top-p", type=float, default=0.95)
    ap.add_argument("--max-tokens", type=int, default=1024)
    ap.add_argument("--concurrency", type=int, default=32)
    a = ap.parse_args()
    cases = [json.loads(line) for line in Path(a.cases).read_text().splitlines() if line.strip()]
    for w in sorted({c["world"] for c in cases if c.get("kind") == "sql"}):
        worldlib.connect(w, worldlib.SEEDS[0])  # build the databases before the threads start
    chat = make_chat(a.endpoint, a.model, a.max_tokens, a.top_p)
    jobs = [(c, s) for c in cases for s in range(a.n)]
    with ThreadPoolExecutor(a.concurrency) as ex:
        done = list(ex.map(lambda j: run_case(chat, j[0], j[1], a.k, a.repairs, a.temperature), jobs))
    out = Path(a.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text("".join(json.dumps(r, ensure_ascii=False) + "\n" for r, _ in done))
    (out.parent / "trace.jsonl").write_text("".join(json.dumps(t, ensure_ascii=False) + "\n" for _, t in done))
    sql = [t for _, t in done if "candidates" in t]
    print(json.dumps({"outputs": len(done), "sql": len(sql), "repaired": sum(any(len(c) > 1 for c in t["candidates"]) for t in sql),
                      "no_valid_candidate": sum(t["chosen"] is None for t in sql), "k": a.k, "repairs": a.repairs, "temperature": a.temperature}))


if __name__ == "__main__":
    main()