"""Spend ledger for the training pod (pure Python; runs on this host). Run as ``python -m training.ledger``.

Classifier-plan §3.4. The pod bills from creation until it is terminated, whether or not a stage runs, so spend is
wall-clock hours since ``started_at`` times the hourly price (an argument: the price checked at rent time) plus
disk. Guards: cap **$50** per training round; no new stage once the projected spend after it passes **$35**
(replan); at **$45** everything stops. Smoke replan rule: measured training throughput below 1,900 tok/s (half the
§4.5 assumption) or a projection above $35 -> drop the screen (use B-last), then 1.5 epochs.

  start  --ledger L --price 2.69 --started-at 2026-11-20T09:00:00+09:00 [--disk-gb 150]
  check  --ledger L --stage NAME --hours H     exit 0 = go, 1 = no new stage (replan line), 2 = stop line
  record --ledger L --stage NAME --event start|end|note [--note TEXT]
  status --ledger L
  replan --ledger L --tokens-per-second X --train-tokens N [--epochs 2] [--screen-hours 0.8] [--eval-hours 0.8]
"""
from __future__ import annotations

import argparse
import json
import os
from datetime import datetime, timezone
from pathlib import Path

VERSION = "ledger-v1"
BUDGET = 50.0
NO_NEW_STAGE = 35.0
STOP = 45.0
MIN_TOKENS_PER_SECOND = 1900.0
REPLAN_EPOCHS = 1.5
DISK_PRICE_GB_MONTH = 0.10
HOURS_PER_MONTH = 730.0
DEFAULT_LEDGER = Path.home() / "naver_psy_models" / "ledger.json"


def parse_time(text: str) -> datetime:
    value = datetime.fromisoformat(text.replace("Z", "+00:00"))
    if value.tzinfo is None:
        raise ValueError("times must carry a UTC offset: " + text)
    return value.astimezone(timezone.utc)


def now() -> datetime:
    return datetime.now(timezone.utc)


def rate(ledger: dict) -> float:
    """$ per wall-clock hour: GPU plus disk."""
    return ledger["price_per_hour"] + ledger["disk_gb"] * DISK_PRICE_GB_MONTH / HOURS_PER_MONTH


def spent(ledger: dict, at: datetime) -> dict:
    hours = max(0.0, (at - parse_time(ledger["started_at"])).total_seconds() / 3600)
    gpu = hours * ledger["price_per_hour"]
    disk = hours * ledger["disk_gb"] * DISK_PRICE_GB_MONTH / HOURS_PER_MONTH
    return {"hours": round(hours, 3), "gpu": round(gpu, 2), "disk": round(disk, 2), "total": round(gpu + disk, 2)}


def check(ledger: dict, stage_hours: float, at: datetime) -> dict:
    """Go / no new stage / stop for a stage of the given estimated length."""
    s = spent(ledger, at)
    projected = round(s["total"] + stage_hours * rate(ledger), 2)
    if s["total"] >= STOP:
        verdict, code = "stop: spend is at the $45 stop line; stop all stages, sync down and terminate", 2
    elif projected > NO_NEW_STAGE:
        verdict, code = "no new stage: projected spend passes the $35 replan line", 1
    else:
        verdict, code = "go", 0
    return {"spent": s, "stage_hours": stage_hours, "projected": projected, "verdict": verdict, "code": code,
            "lines": {"no_new_stage": NO_NEW_STAGE, "stop": STOP, "budget": BUDGET}}


def train_hours(train_tokens: int, tokens_per_second: float, epochs: float) -> float:
    return train_tokens * epochs / tokens_per_second / 3600


def replan(ledger: dict, tokens_per_second: float, train_tokens: int, at: datetime, epochs: float = 2.0,
           screen_hours: float = 0.8, eval_hours: float = 0.8) -> dict:
    """The smoke rule: projection of the remaining plan, and the reduced plan when it is needed."""
    s = spent(ledger, at)
    r = rate(ledger)
    final = train_hours(train_tokens, tokens_per_second, epochs)
    full = round(s["total"] + (screen_hours + final + eval_hours) * r, 2)
    slow = tokens_per_second < MIN_TOKENS_PER_SECOND
    reasons = []
    if slow:
        reasons.append(f"throughput {tokens_per_second:.0f} tok/s < {MIN_TOKENS_PER_SECOND:.0f}")
    if full > NO_NEW_STAGE:
        reasons.append(f"projection ${full:.2f} > ${NO_NEW_STAGE:.0f}")
    out = {"spent": s, "tokens_per_second": tokens_per_second, "train_tokens": train_tokens,
           "plan": {"screen": True, "epochs": epochs, "final_hours": round(final, 2), "projected": full},
           "replan": bool(reasons), "reasons": reasons}
    if reasons:
        reduced_final = train_hours(train_tokens, tokens_per_second, REPLAN_EPOCHS)
        reduced = round(s["total"] + (reduced_final + eval_hours) * r, 2)
        out["reduced_plan"] = {"screen": False, "arm": "B-last", "epochs": REPLAN_EPOCHS,
                               "final_hours": round(reduced_final, 2), "projected": reduced,
                               "fits": reduced <= NO_NEW_STAGE}
    return out


def load(path) -> dict:
    return json.loads(Path(path).read_text())


def save(path, ledger: dict) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_name(path.name + ".tmp")
    tmp.write_text(json.dumps(ledger, indent=1) + "\n")
    os.chmod(tmp, 0o600)
    tmp.replace(path)


def new_ledger(price: float, started_at: str, disk_gb: float, at: datetime) -> dict:
    if price <= 0:
        raise ValueError("price must be positive")
    parse_time(started_at)
    return {"version": VERSION, "started_at": started_at, "price_per_hour": price, "disk_gb": disk_gb,
            "events": [{"at": at.isoformat(), "event": "ledger-start"}]}


def main(argv=None) -> int:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    sub = p.add_subparsers(dest="command", required=True)
    s = sub.add_parser("start")
    s.add_argument("--price", type=float, required=True, help="GPU price in $ per hour (checked at rent time)")
    s.add_argument("--started-at", required=True, help="When the pod was created (ISO time with offset)")
    s.add_argument("--disk-gb", type=float, default=150.0, help="Container disk plus network volume")
    c = sub.add_parser("check")
    c.add_argument("--stage", required=True)
    c.add_argument("--hours", type=float, required=True, help="Estimated hours of the next stage")
    r = sub.add_parser("record")
    r.add_argument("--stage", required=True)
    r.add_argument("--event", required=True, choices=["start", "end", "note"])
    r.add_argument("--note", default="")
    sub.add_parser("status")
    q = sub.add_parser("replan")
    q.add_argument("--tokens-per-second", type=float, required=True)
    q.add_argument("--train-tokens", type=int, required=True)
    q.add_argument("--epochs", type=float, default=2.0)
    q.add_argument("--screen-hours", type=float, default=0.8)
    q.add_argument("--eval-hours", type=float, default=0.8)
    for sp in (s, c, r, q, sub.choices["status"]):
        sp.add_argument("--ledger", default=str(DEFAULT_LEDGER))
    args = p.parse_args(argv)
    at = now()
    if args.command == "start":
        if Path(args.ledger).exists():
            raise SystemExit("refusing to overwrite " + args.ledger)
        save(args.ledger, new_ledger(args.price, args.started_at, args.disk_gb, at))
        print(json.dumps(spent(load(args.ledger), at)))
        return 0
    ledger = load(args.ledger)
    if args.command == "check":
        result = check(ledger, args.hours, at)
        ledger["events"].append({"at": at.isoformat(), "event": "check", "stage": args.stage, **result})
        save(args.ledger, ledger)
        print(json.dumps(result))
        return result["code"]
    if args.command == "replan":
        result = replan(ledger, args.tokens_per_second, args.train_tokens, at, args.epochs, args.screen_hours,
                        args.eval_hours)
        ledger["events"].append({"at": at.isoformat(), "event": "replan", **result})
        save(args.ledger, ledger)
        print(json.dumps(result))
        return 1 if result["replan"] else 0
    if args.command == "record":
        ledger["events"].append({"at": at.isoformat(), "event": args.event, "stage": args.stage, "note": args.note,
                                 "spent": spent(ledger, at)})
        save(args.ledger, ledger)
    print(json.dumps({"spent": spent(ledger, at), "events": len(ledger["events"]),
                      "price_per_hour": ledger["price_per_hour"]}))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())