#!/usr/bin/env python3
"""Tiny durable loop harness primitives.

Loop engineering is just: pick next task, run one bounded step, write a receipt.
No dashboard required until this boring JSONL spine is proven useful.
"""

from __future__ import annotations

import argparse
import json
import os
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from typing import Any, Iterable, Optional


@dataclass
class LoopTask:
    id: str
    title: str
    status: str = "todo"


@dataclass
class LoopEvent:
    run_id: str
    type: str
    data: dict[str, Any]
    timestamp: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())


@dataclass
class CostBudget:
    max_tokens: int
    used_tokens: int = 0

    def allow(self, estimate_tokens: int) -> bool:
        return self.used_tokens + estimate_tokens <= self.max_tokens

    def record(self, tokens: int) -> None:
        self.used_tokens += tokens


def load_tasks(path: str) -> list[LoopTask]:
    with open(path, "r", encoding="utf-8") as f:
        raw = json.load(f)
    return [LoopTask(**item) for item in raw]


def save_tasks(path: str, tasks: Iterable[LoopTask]) -> None:
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump([asdict(task) for task in tasks], f, indent=2)
        f.write("\n")


def next_task(tasks: Iterable[LoopTask]) -> Optional[LoopTask]:
    return next((task for task in tasks if task.status == "todo"), None)


def append_event(path: str, event: LoopEvent) -> None:
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    with open(path, "a", encoding="utf-8") as f:
        f.write(json.dumps(asdict(event), ensure_ascii=False) + "\n")


def read_events(path: str) -> list[LoopEvent]:
    if not os.path.exists(path):
        return []
    with open(path, "r", encoding="utf-8") as f:
        return [LoopEvent(**json.loads(line)) for line in f if line.strip()]


def run_one_cycle(tasks_path: str, events_path: str, run_id: str, budget: CostBudget, estimate_tokens: int = 0) -> dict[str, Any]:
    """Claim one todo task if budget allows; callers do the actual agent work."""
    if not budget.allow(estimate_tokens):
        event = LoopEvent(run_id, "blocked", {"reason": "token_budget", "estimate_tokens": estimate_tokens, "budget": asdict(budget)})
        append_event(events_path, event)
        return {"status": "blocked", "reason": "token_budget"}

    tasks = load_tasks(tasks_path)
    task = next_task(tasks)
    if task is None:
        event = LoopEvent(run_id, "done", {"reason": "no_todo_tasks"})
        append_event(events_path, event)
        return {"status": "done"}

    task.status = "claimed"
    budget.record(estimate_tokens)
    save_tasks(tasks_path, tasks)
    event = LoopEvent(run_id, "claimed", {"task": asdict(task), "estimate_tokens": estimate_tokens, "budget": asdict(budget)})
    append_event(events_path, event)
    return {"status": "claimed", "task": asdict(task), "budget": asdict(budget)}


def main() -> int:
    parser = argparse.ArgumentParser(description="Run one bounded loop-harness cycle")
    parser.add_argument("--tasks", required=True, help="JSON task list")
    parser.add_argument("--events", required=True, help="JSONL event log")
    parser.add_argument("--run-id", required=True)
    parser.add_argument("--max-tokens", type=int, default=20_000)
    parser.add_argument("--used-tokens", type=int, default=0)
    parser.add_argument("--estimate-tokens", type=int, default=0)
    args = parser.parse_args()

    receipt = run_one_cycle(
        args.tasks,
        args.events,
        args.run_id,
        CostBudget(args.max_tokens, args.used_tokens),
        args.estimate_tokens,
    )
    print(json.dumps(receipt, indent=2))
    return 1 if receipt.get("status") == "blocked" else 0


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