"""
Execution-accuracy + latency benchmark for the Phi-3 Text-to-SQL GGUF models.

For each held-out test question we:
  1. send the EXACT production prompt (schema + question) to a llama-server
     instance serving one GGUF quantization,
  2. extract the SQL from the model output,
  3. execute predicted SQL and gold SQL against the real SQLite DB,
  4. score execution-match (order-insensitive row-set equality).

Reports per-model: execution accuracy, valid-SQL rate, exact-string match,
and decode latency stats. Genuine numbers only.
"""
import json, os, re, sys, time, subprocess, urllib.request, urllib.error, sqlite3

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DB = os.path.join(ROOT, "data", "company_sales.db")
TEST = os.path.join(ROOT, "data", "test_dataset.jsonl")
SERVER = os.path.join(ROOT, "llama-b9637-bin-win-cpu-x64", "llama-server.exe")
PORT = 8087
THREADS = 14


def load_tests():
    rows = []
    with open(TEST, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            msgs = json.loads(line)["messages"]
            user = next(m["content"] for m in msgs if m["role"] == "user")
            gold = next(m["content"] for m in msgs if m["role"] == "assistant")
            rows.append((user, gold))
    return rows


def clean_sql(raw):
    if "<|assistant|>" in raw:
        raw = raw.split("<|assistant|>")[-1]
    raw = re.sub(r"```sql\s*", "", raw, flags=re.IGNORECASE)
    raw = re.sub(r"```", "", raw)
    raw = re.sub(r"<\|.*?\|>", "", raw).strip()
    m = re.search(r"(SELECT\s+.*?;)", raw, re.DOTALL | re.IGNORECASE)
    if m:
        return m.group(1).strip()
    return raw.strip()


def run_sql(sql):
    try:
        conn = sqlite3.connect(DB)
        cur = conn.execute(sql)
        rows = cur.fetchall()
        conn.close()
        return frozenset(map(tuple, rows)), None
    except Exception as e:
        return None, str(e)


def wait_health(timeout=120):
    t0 = time.time()
    while time.time() - t0 < timeout:
        try:
            with urllib.request.urlopen(f"http://127.0.0.1:{PORT}/health", timeout=2) as r:
                if json.loads(r.read()).get("status") == "ok":
                    return True
        except Exception:
            time.sleep(1)
    return False


def query(user_prompt):
    body = json.dumps({
        "messages": [{"role": "user", "content": user_prompt}],
        "temperature": 0.0,
        "max_tokens": 150,
        "cache_prompt": False,
    }).encode()
    req = urllib.request.Request(
        f"http://127.0.0.1:{PORT}/v1/chat/completions",
        data=body, headers={"Content-Type": "application/json"})
    t0 = time.perf_counter()
    with urllib.request.urlopen(req, timeout=180) as r:
        resp = json.loads(r.read())
    dt = time.perf_counter() - t0
    text = resp["choices"][0]["message"]["content"]
    usage = resp.get("usage", {})
    ct = usage.get("completion_tokens", 0)
    return text, dt, ct


def bench(model_path):
    name = os.path.basename(model_path)
    proc = subprocess.Popen(
        [SERVER, "-m", model_path, "-t", str(THREADS), "--port", str(PORT),
         "-c", "2048", "-ngl", "0", "--no-webui"],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    try:
        if not wait_health():
            print(f"[{name}] server failed to start"); return None
        tests = load_tests()
        n = len(tests)
        exec_ok = valid_sql = exact = 0
        lat, toks = [], []
        details = []
        for i, (user, gold) in enumerate(tests, 1):
            text, dt, ct = query(user)
            pred = clean_sql(text)
            lat.append(dt); toks.append(ct)
            pset, perr = run_sql(pred)
            gset, gerr = run_sql(gold)
            is_valid = perr is None
            is_exec = is_valid and gerr is None and pset == gset
            is_exact = pred.strip().rstrip(";").lower() == gold.strip().rstrip(";").lower()
            valid_sql += is_valid; exec_ok += is_exec; exact += is_exact
            details.append({"q": user.split("Question:")[-1].strip(),
                            "pred": pred, "gold": gold,
                            "valid": is_valid, "exec_match": is_exec,
                            "gold_err": gerr})
        tg = [c / l for c, l in zip(toks, lat) if l > 0 and c]
        res = {
            "model": name,
            "n": n,
            "execution_accuracy_pct": round(100 * exec_ok / n, 1),
            "valid_sql_pct": round(100 * valid_sql / n, 1),
            "exact_match_pct": round(100 * exact / n, 1),
            "avg_latency_s": round(sum(lat) / n, 2),
            "avg_decode_tok_s": round(sum(tg) / len(tg), 1) if tg else 0,
            "details": details,
        }
        return res
    finally:
        proc.terminate()
        try:
            proc.wait(timeout=15)
        except Exception:
            proc.kill()
        time.sleep(2)


if __name__ == "__main__":
    models = sys.argv[1:] or [
        os.path.join(ROOT, "models", "phi3-text-to-sql-Q5_K_M.gguf"),
        os.path.join(ROOT, "models", "phi3-text-to-sql-Q4_K_M.gguf"),
    ]
    all_res = []
    for m in models:
        r = bench(m)
        if r:
            all_res.append(r)
            print(f"\n=== {r['model']} ===")
            print(f"  execution accuracy : {r['execution_accuracy_pct']}%  ({r['n']} questions)")
            print(f"  valid SQL rate     : {r['valid_sql_pct']}%")
            print(f"  exact-string match : {r['exact_match_pct']}%")
            print(f"  avg latency        : {r['avg_latency_s']} s")
            print(f"  avg decode speed   : {r['avg_decode_tok_s']} tok/s")
    with open(os.path.join(ROOT, "scripts", "eval_results.json"), "w") as f:
        json.dump(all_res, f, indent=2)
    print("\nSaved scripts/eval_results.json")