"""v1.1 WP-A1: real turns, and what the server said it read.

    AIDND_TEST_ENDPOINT=https://<host>:<port>/v1 AIDND_TEST_MODEL=<model> \\
      .venv/bin/python -m tools.v11_window_accounting \\
        --bundle "$HOME/m11-evidence/m04-final/bundle.json" --turns 4 \\
        --out "$HOME/v11-evidence/a1-accounting/<label>"

Run from `backend/`. The evidence that matters is at the edge of the window, and
a new campaign takes dozens of turns to reach it. So this imports a long
campaign, by default the v1 evidence run's 207-action bundle, and every turn is
assembled against a full window from the first. A real narrator is then asked
for `--turns` turns, and each one is written out with:

    configured budget, verified window, and where the window came from
    the application's estimate of what it sent (its count, plus the text the
    provider adds)
    the server's own prompt-token count, from the usage it reported
    the reply allocation and the safety reserve
    the observed margin: window - reply allocation - the server's count
    the accounting status: fits, exceeded, truncation_suspected or unknown

The first turn on a cold model cannot verify the window: `/api/ps` knows nothing
until the model is loaded. That is M11's behaviour, and the row says so rather
than hiding the turn.

The database lives in `--out`, not in `/tmp`, so the stored snapshots behind
every row can be read again. The endpoint is read from the environment and is
never written into a committed file.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from pathlib import Path

ENDPOINT = os.environ.get("AIDND_TEST_ENDPOINT", "")
MODEL = os.environ.get("AIDND_TEST_MODEL", "")

TURNS = [
    "I look around carefully and take stock of where I am.",
    "I ask the nearest person what has happened since I was last here.",
    "I check what I am carrying.",
    "I move on towards the place I meant to reach.",
    "I wait and listen.",
    "I say, \"Tell me the part you left out.\"",
]


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument("--bundle", default="",
                        help="a campaign bundle to import, so turns start at a full window")
    parser.add_argument("--turns", type=int, default=4)
    parser.add_argument("--budget", type=int, default=16384)
    parser.add_argument("--max-output", type=int, default=500)
    parser.add_argument("--timeout", type=int, default=900)
    parser.add_argument("--out", required=True)
    parser.add_argument("--unload-first", action="store_true",
                        help="ask the configured server to unload the model before turn 1, "
                             "so the first turn starts cold (the A1 corrective test)")
    args = parser.parse_args()

    if not (ENDPOINT and MODEL):
        print("set AIDND_TEST_ENDPOINT and AIDND_TEST_MODEL")
        return 2

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    db_path = out / "accounting.db"
    if db_path.exists():
        print(f"{db_path} exists; choose a new --out")
        return 2
    os.environ["AIDND_DB_PATH"] = str(db_path)
    os.environ.pop("AIDND_DATABASE_URL", None)
    os.environ.pop("DATABASE_URL", None)

    from fastapi import Depends
    from fastapi.testclient import TestClient
    from sqlalchemy.orm import undefer

    from app import auth, limits, models
    from app.database import Base, SessionLocal, engine, get_db
    from app.main import app

    limits.check_row_cap = lambda *a, **k: None
    Base.metadata.create_all(bind=engine)
    with SessionLocal() as db:
        user = models.User(is_guest=False, email="accounting@example.com")
        db.add(user)
        db.flush()
        db.add(models.Settings(
            user_id=user.id, model=MODEL, endpoint_url=ENDPOINT, embedding_model="",
            context_token_budget=args.budget, max_output_tokens=args.max_output,
            model_timeout_seconds=args.timeout,
        ))
        db.commit()
        user_id = user.id
    app.dependency_overrides[auth.get_current_user] = (
        lambda db=Depends(get_db): db.get(models.User, user_id)
    )
    client = TestClient(app)

    if args.bundle:
        bundle = json.loads(Path(args.bundle).read_text())
        # The evidence campaign had its memory bank and auto-summarise on. Here
        # they would only add post-turn model calls between the measured turns,
        # on the same host, and fail noisily as the in-process client closes.
        # This tool measures the turn's own prompt, which neither changes.
        bundle["memoryBankEnabled"] = False
        bundle["autoSummarize"] = False
        imported = client.post("/api/adventures/import", json=bundle)
        imported.raise_for_status()
        adv = imported.json()["id"]
    else:
        created = client.post("/api/adventures", json={
            "title": "Window accounting", "opening": "A quiet road at dusk."})
        created.raise_for_status()
        adv = created.json()["id"]

    if args.unload_first:
        # The configured endpoint only, under the same policy and TLS trust as a turn.
        import asyncio
        import httpx
        from app import contextwindow, endpoints, tlstrust
        reason = endpoints.rejection_reason(ENDPOINT)
        if reason:
            print(f"endpoint refused: {reason}")
            return 2
        base = contextwindow.native_base(ENDPOINT)
        with httpx.Client(verify=tlstrust.ssl_context(), timeout=120) as http:
            unloaded = http.post(f"{base}/api/generate", json={"model": MODEL, "keep_alive": 0})
            resident = [m.get("name") for m in http.get(f"{base}/api/ps").json().get("models", [])]
        contextwindow.cache_clear()
        print(f"unload: HTTP {unloaded.status_code} {unloaded.text[:120]} | resident now: {resident}")

    rows: list[dict] = []
    timeline = (out / "turns.jsonl").open("a")
    print(f"model {MODEL}, budget {args.budget}, reply {args.max_output}")
    print(f"{'#':>2} {'status':22} {'window':>13} {'estimate':>8} {'server':>7} "
          f"{'reserve':>7} {'margin':>7} {'sec':>5}")
    for index in range(args.turns):
        text = TURNS[index % len(TURNS)]
        started = time.monotonic()
        response = client.post(f"/api/adventures/{adv}/actions",
                               json={"type": "do", "text": text})
        seconds = round(time.monotonic() - started, 1)
        error = None
        if response.status_code != 200 or '"type": "error"' in response.text:
            error = response.text[-400:]
        with SessionLocal() as db:
            action = (
                db.query(models.Action)
                .filter(models.Action.adventure_id == adv, models.Action.type == "ai")
                .options(undefer(models.Action.context_snapshot))
                .order_by(models.Action.id.desc()).first()
            )
            snapshot = (action.context_snapshot or {}) if action else {}
        tokens = snapshot.get("tokens") or {}
        window = snapshot.get("window") or {}
        accounting = snapshot.get("accounting") or {}
        row = {
            "turn": index + 1,
            "action_id": action.id if action else None,
            "seconds": seconds,
            "error": error,
            "configured_budget": tokens.get("configured_budget"),
            "effective_budget": tokens.get("budget"),
            "window_verified": window.get("verified"),
            "window_tokens": window.get("tokens"),
            "window_source": window.get("source"),
            "preflight_attempted": (window.get("preflight") or {}).get("attempted"),
            "preflight_loaded": (window.get("preflight") or {}).get("loaded"),
            "preflight_verified_before": (window.get("preflight") or {}).get("verified_before"),
            "preflight_verified_after": (window.get("preflight") or {}).get("verified_after"),
            "preflight_detail": (window.get("preflight") or {}).get("detail"),
            "app_prompt_tokens": tokens.get("total"),
            "transport_tokens": tokens.get("transport"),
            "app_estimate": tokens.get("estimate"),
            "output_reserve": tokens.get("output_reserve"),
            "safety_reserve": tokens.get("safety_reserve"),
            "server_prompt_tokens": accounting.get("server_prompt_tokens"),
            "difference": accounting.get("difference"),
            "observed_margin": accounting.get("observed_margin"),
            "status": accounting.get("status"),
            "history_included": (snapshot.get("history") or {}).get("included"),
            "history_total": (snapshot.get("history") or {}).get("total"),
        }
        rows.append(row)
        timeline.write(json.dumps(row) + "\n")
        timeline.flush()
        print(f"{row['turn']:>2} {str(row['status'] if not error else 'ERROR'):22} "
              f"{str(row['window_tokens'])+('v' if row['window_verified'] else '?'):>13} "
              f"{str(row['app_estimate']):>8} {str(row['server_prompt_tokens']):>7} "
              f"{str(row['safety_reserve']):>7} {str(row['observed_margin']):>7} {seconds:>5}")
        if error:
            print(f"   error: {error[:200]}")
    timeline.close()
    (out / "summary.json").write_text(json.dumps({
        "model": MODEL, "budget": args.budget, "max_output_tokens": args.max_output,
        "bundle": args.bundle, "rows": rows,
    }, indent=2))
    app.dependency_overrides.clear()
    return 0 if all(r["error"] is None for r in rows) else 1


if __name__ == "__main__":
    sys.exit(main())