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

The pod bills from the moment the user creates it until it is stopped or terminated, whether or not a stage
is running, so spend is wall-clock time since ``started_at`` times the hourly price, plus disk. Guards
(spec section 6): no new stage may start if the projected spend after it passes NO_NEW_STAGE, and at STOP
everything stops and the user is told. Only the user stops or terminates the pod (user decision).

  start  --ledger L --price 3.99 --started-at 2026-10-17T09:00:00Z [--disk-gb 150]
  check  --ledger L --stage NAME --hours H     exit 0 = go, 1 = over the no-new-stage line, 2 = at the stop line
  record --ledger L --stage NAME --event start|end|note [--note TEXT]
  status --ledger L
"""
from __future__ import annotations

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

BUDGET = 50.0
NO_NEW_STAGE = 35.0
STOP = 45.0
DISK_PRICE_GB_MONTH = 0.10
HOURS_PER_MONTH = 730.0


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 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:
    s = spent(ledger, at)
    rate = ledger['price_per_hour'] + ledger['disk_gb'] * DISK_PRICE_GB_MONTH / HOURS_PER_MONTH
    projected = round(s['total'] + stage_hours * rate, 2)
    if s['total'] >= STOP:
        verdict, code = 'stop: spend is at the stop line; stop all stages and tell the user', 2
    elif projected > NO_NEW_STAGE:
        verdict, code = 'no new stage: projected spend passes the no-new-stage 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 load(path) -> dict:
    return json.loads(Path(path).read_text())


def save(path, ledger: dict) -> None:
    path = Path(path)
    tmp = path.with_suffix('.tmp')
    tmp.write_text(json.dumps(ledger, indent=1) + '\n')
    tmp.replace(path)


def main(argv=None):
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    sub = p.add_subparsers(dest='command', required=True)
    s = sub.add_parser('start')
    s.add_argument('--ledger', required=True)
    s.add_argument('--price', type=float, required=True, help='GPU price in $ per hour')
    s.add_argument('--started-at', required=True, help='When the user created the pod (ISO time with offset)')
    s.add_argument('--disk-gb', type=float, default=150.0, help='Container disk plus volume')
    c = sub.add_parser('check')
    c.add_argument('--ledger', required=True)
    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('--ledger', required=True)
    r.add_argument('--stage', required=True)
    r.add_argument('--event', required=True, choices=['start', 'end', 'note'])
    r.add_argument('--note', default='')
    t = sub.add_parser('status')
    t.add_argument('--ledger', required=True)
    args = p.parse_args(argv)
    at = now()
    if args.command == 'start':
        if Path(args.ledger).exists():
            raise SystemExit('Refusing to overwrite ' + args.ledger)
        parse_time(args.started_at)
        save(args.ledger, {'version': 'ledger-v1', 'started_at': args.started_at, 'price_per_hour': args.price,
                           'disk_gb': args.disk_gb, 'events': [{'at': at.isoformat(), 'event': 'ledger-start'}]})
        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 == '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'])}))
    return 0


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