"""loop.py — Karpathy-style autoresearch loop for jepa-fx-risk.

Agent (on iguana/berget — NOT koala, whose GPU is reserved for train.py) reads
program.md + train.py + STATUS.md, proposes ONE change to train.py, we run it,
keep if val_vol_r2 improved else git-revert. Appends per-iter record to STATUS.md.

    LITELLM_KEY=xxx python loop.py [--iters N] [--model MODEL] [--run-dir runs/rq-04]

Env:
  LITELLM_KEY      — LiteLLM master key (required)
  LITELLM_BASE     — default http://localhost:30401/v1
  LOOP_MODEL       — default berget/gemma4-31b (non-thinking; iguana/berget only)
  LOOP_ITERS       — default 3
  TRAIN_TIMEOUT    — seconds per train.py run, default 120
  NTFY_URL         — optional: POST crash/stall alerts here (e.g. ntfy.sh/<topic>)
"""
import argparse
import json
import os
import subprocess
import sys
import time
import textwrap
from pathlib import Path

import urllib.request

LITELLM_BASE  = os.environ.get("LITELLM_BASE", "http://localhost:30401/v1")
LITELLM_KEY   = os.environ.get("LITELLM_KEY", "")
LOOP_MODEL    = os.environ.get("LOOP_MODEL", "berget/gemma4-31b")
LOOP_ITERS    = int(os.environ.get("LOOP_ITERS", "3"))
TRAIN_TIMEOUT = int(os.environ.get("TRAIN_TIMEOUT", "120"))
NTFY_URL      = os.environ.get("NTFY_URL", "")

# Resolved by main() once --run-dir is parsed.
RUN_DIR      = Path(".")
STATUS_MD    = Path("STATUS.md")
METRICS_JSON = Path("metrics.json")
TRAIN_PY     = Path("train.py")
HEARTBEAT    = Path("HEARTBEAT")

AGENT_SYSTEM = textwrap.dedent("""\
    You are the autoresearch agent for jepa-fx-risk. Your job: propose ONE small,
    targeted change to train.py to improve val_vol_r2 (OOS R² predicting 1-day
    realized vol from frozen embeddings). Higher is better.

    Rules:
    - Return ONLY the full new content of train.py — nothing else, no explanation,
      no markdown fence. Raw Python only.
    - Change ONE thing at a time (one knob, one structural idea).
    - Do NOT touch prepare_data.py, loop.py, or the data pipeline — only train.py.
    - Do NOT add new data sources or new files.
    - The metric is computed externally from your frozen embeddings; trust it.
""")


def read_file(p: Path) -> str:
    return p.read_text() if p.exists() else ""


def gpu_snapshot() -> str:
    try:
        out = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=utilization.gpu,memory.used,memory.total,temperature.gpu",
             "--format=csv,noheader,nounits"], timeout=5, text=True
        ).strip()
        util, mem_used, mem_total, temp = [x.strip() for x in out.split(",")]
        return "gpu=%s%% vram=%s/%sMiB temp=%s°C" % (util, mem_used, mem_total, temp)
    except Exception:
        return "gpu=N/A"


def read_metric() -> "float | None":
    if not METRICS_JSON.exists():
        return None
    try:
        return float(json.loads(METRICS_JSON.read_text())["val_vol_r2"])
    except Exception:
        return None


def run_train() -> "tuple[float | None, float, str]":
    """Run train.py from project root with METRICS_OUT pointing into the run dir."""
    t0 = time.time()
    env = dict(os.environ)
    env["METRICS_OUT"] = str(METRICS_JSON.resolve())
    # train.py is copied into the run dir, so sys.path[0] is that run dir — which
    # has no scripts/. Put the project root (where loop.py + scripts/ live) on
    # PYTHONPATH so train.py's `from scripts.var_breach import ...` resolves.
    env["PYTHONPATH"] = str(Path(__file__).resolve().parent) + os.pathsep + env.get("PYTHONPATH", "")
    try:
        r = subprocess.run(
            [sys.executable, str(TRAIN_PY.resolve())],
            capture_output=True, text=True, timeout=TRAIN_TIMEOUT, env=env,
        )
        elapsed = time.time() - t0
        if r.returncode != 0:
            return None, elapsed, (r.stderr or r.stdout)[-300:]
        metric = read_metric()
        return metric, elapsed, ""
    except subprocess.TimeoutExpired:
        return None, TRAIN_TIMEOUT, "TIMEOUT"


def call_agent(iteration: int, best_so_far: "float | None") -> str:
    """Ask the LLM agent to edit train.py. Returns new train.py content."""
    context = "\n\n".join([
        "# program.md\n" + read_file(RUN_DIR / "program.md"),
        "# train.py (current)\n" + read_file(TRAIN_PY),
        "# STATUS.md (history)\n" + read_file(STATUS_MD)[-2000:],
        "# metrics.json (last run)\n" + read_file(METRICS_JSON),
        "Iteration %d. Best val_vol_r2 so far: %s. Improve it." % (
            iteration, "%.4f" % best_so_far if best_so_far is not None else "none yet"
        ),
    ])
    payload = json.dumps({
        "model": LOOP_MODEL,
        "messages": [
            {"role": "system", "content": AGENT_SYSTEM},
            {"role": "user",   "content": context},
        ],
        "temperature": 0.7,
        "max_tokens": 4096,
    }).encode()
    req = urllib.request.Request(
        LITELLM_BASE + "/chat/completions",
        data=payload,
        headers={"Authorization": "Bearer " + LITELLM_KEY,
                 "Content-Type": "application/json"},
        method="POST",
    )
    resp = urllib.request.urlopen(req, timeout=60)
    data = json.load(resp)
    return data["choices"][0]["message"]["content"]


def revert_train(original_content: str):
    TRAIN_PY.write_text(original_content)


def append_status(line: str):
    with open(STATUS_MD, "a") as f:
        f.write(line + "\n")


def write_heartbeat(iteration: int, status: str = "alive"):
    """Update HEARTBEAT so watchdogs can detect stalls."""
    HEARTBEAT.write_text("%s  iter=%d  ts=%.0f\n" % (status, iteration, time.time()))


def ntfy(msg: str):
    """POST an alert to NTFY_URL (best-effort; silently ignored on any error)."""
    if not NTFY_URL:
        return
    try:
        req = urllib.request.Request(
            NTFY_URL, data=msg.encode(), method="POST",
            headers={"Content-Type": "text/plain"},
        )
        urllib.request.urlopen(req, timeout=5)
    except Exception:
        pass


def main():
    global RUN_DIR, STATUS_MD, METRICS_JSON, TRAIN_PY, HEARTBEAT

    parser = argparse.ArgumentParser()
    parser.add_argument("--iters", type=int, default=LOOP_ITERS)
    parser.add_argument("--model", default=LOOP_MODEL)
    parser.add_argument(
        "--run-dir", default=None,
        help="run dir scaffolded by autoresearch_start.py; "
             "STATUS.md, metrics.json, HEARTBEAT, and train.py live here",
    )
    args = parser.parse_args()

    loop_iters = args.iters
    loop_model = args.model

    if args.run_dir:
        RUN_DIR = Path(args.run_dir)
        if not RUN_DIR.is_dir():
            print("ERROR: run dir not found:", RUN_DIR); sys.exit(1)

    STATUS_MD    = RUN_DIR / "STATUS.md"
    METRICS_JSON = RUN_DIR / "metrics.json"
    TRAIN_PY     = RUN_DIR / "train.py"
    HEARTBEAT    = RUN_DIR / "HEARTBEAT"

    if not LITELLM_KEY:
        print("ERROR: set LITELLM_KEY"); sys.exit(1)

    if not STATUS_MD.exists():
        STATUS_MD.write_text(
            "# Autoresearch STATUS\n\n"
            "| iter | val_vol_r2 | delta | action | secs | gpu | change |\n"
            "|------|-----------|-------|--------|------|-----|--------|\n"
        )

    baseline = read_metric()
    if baseline is None:
        print("No metrics.json — running train.py for baseline...")
        m, secs, err = run_train()
        if m is None:
            msg = "Baseline run failed: " + err
            print(msg)
            ntfy("[jepa-fx-risk] loop CRASH — " + msg)
            sys.exit(1)
        baseline = m
        print("Baseline: val_vol_r2 = %.4f (%.1fs)" % (baseline, secs))

    best = baseline
    print("Starting loop | model=%s | iters=%d | baseline=%.4f" % (loop_model, loop_iters, best))
    if args.run_dir:
        print("  run-dir:", RUN_DIR)

    iter_index = 0
    try:
        for i in range(1, loop_iters + 1):
            iter_index = i
            write_heartbeat(i, "agent-call")
            print("\n--- iter %d/%d ---" % (i, loop_iters))
            original = TRAIN_PY.read_text()

            print("  calling agent (%s)..." % loop_model)
            t_agent = time.time()
            try:
                new_code = call_agent(i, best)
            except Exception as e:
                msg = str(e)
                print("  agent call failed:", msg)
                append_status("| %d | ERR | — | agent-fail | — | — | %s |" % (i, msg[:60]))
                write_heartbeat(i, "agent-fail")
                ntfy("[jepa-fx-risk] iter %d agent FAIL — %s" % (i, msg[:80]))
                continue
            agent_secs = time.time() - t_agent
            print("  agent replied in %.1fs" % agent_secs)

            # strip accidental markdown fences
            if new_code.strip().startswith("```"):
                lines = new_code.strip().splitlines()
                new_code = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:])

            TRAIN_PY.write_text(new_code)

            write_heartbeat(i, "training")
            gpu = gpu_snapshot()
            print("  running train.py [%s]..." % gpu)
            metric, secs, err = run_train()

            if metric is None:
                print("  train.py FAILED — reverting. err:", err[:100])
                revert_train(original)
                append_status("| %d | FAIL | — | revert | %.0fs | %s | run error |" % (i, secs, gpu))
                write_heartbeat(i, "train-fail")
                ntfy("[jepa-fx-risk] iter %d train FAIL — %s" % (i, err[:80]))
                continue

            delta = metric - best
            if metric > best:
                best = metric
                action = "KEEP"
            else:
                revert_train(original)
                action = "revert"

            summary = "| %d | %.4f | %+.4f | %s | %.0fs | %s | iter%d |" % (
                i, metric, delta, action, secs, gpu, i)
            append_status(summary)
            write_heartbeat(i, "done")
            print("  val_vol_r2=%.4f  delta=%+.4f  action=%s  [%.0fs]" % (metric, delta, action, secs))

    except Exception as e:
        msg = "loop CRASH at iter %d: %s" % (iter_index, e)
        print("FATAL:", msg)
        ntfy("[jepa-fx-risk] " + msg)
        raise

    print("\nDone. Best val_vol_r2 = %.4f (baseline was %.4f, delta %+.4f)" % (best, baseline, best - baseline))
    print("STATUS.md updated.")
    write_heartbeat(loop_iters, "done")
    ntfy("[jepa-fx-risk] loop done. best val_vol_r2=%.4f (delta %+.4f)" % (best, best - baseline))


if __name__ == "__main__":
    main()