#!/usr/bin/env python3
"""Mini-only drafting guardrail.

Use this to create a local proof trail before any task that must be attributed to
GPT 5.4 mini. It does not spawn subagents by itself; instead it records the
spawn receipt, stores the exact child-produced draft, and verifies the artifact
before a send step.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
import shutil
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

BASE = Path("/home/isthekid/.openclaw/workspace/state/mini_guardrail")
RECORDS = BASE / "records"
DRAFTS = BASE / "drafts"
MINI_ALIASES = {"gpt-mini", "openai/gpt-5.4-mini"}


def now_iso() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()


def ensure_dirs() -> None:
    RECORDS.mkdir(parents=True, exist_ok=True)
    DRAFTS.mkdir(parents=True, exist_ok=True)


def slugify(value: str) -> str:
    value = value.strip().lower()
    value = re.sub(r"[^a-z0-9._-]+", "-", value)
    value = re.sub(r"-+", "-", value).strip("-._")
    if not value:
        raise SystemExit("slug cannot be empty")
    return value


def record_path(slug: str) -> Path:
    return RECORDS / f"{slug}.json"


def load_record(slug: str) -> dict[str, Any]:
    path = record_path(slug)
    if not path.exists():
        raise SystemExit(f"record not found: {path}")
    return json.loads(path.read_text())


def save_record(slug: str, data: dict[str, Any]) -> None:
    record_path(slug).write_text(json.dumps(data, indent=2, ensure_ascii=False) + "\n")


def sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            h.update(chunk)
    return h.hexdigest()


def is_mini(model: str | None) -> bool:
    return (model or "").strip() in MINI_ALIASES


def print_json(payload: dict[str, Any], ok_exit: bool = True) -> None:
    print(json.dumps(payload, indent=2, ensure_ascii=False))
    if not ok_exit:
        sys.exit(1)


def cmd_init(args: argparse.Namespace) -> None:
    ensure_dirs()
    slug = slugify(args.slug)
    path = record_path(slug)
    if path.exists() and not args.force:
        raise SystemExit(f"record already exists: {path} (use --force to overwrite)")

    require_model = args.require_model.strip()
    data = {
        "slug": slug,
        "taskKind": args.task_kind,
        "requestedBy": args.requested_by,
        "requireModel": require_model,
        "status": "initialized",
        "createdAt": now_iso(),
        "notes": [args.note] if args.note else [],
    }
    save_record(slug, data)
    print_json({"ok": True, "record": str(path), "slug": slug, "status": data["status"]})


def cmd_claim_child(args: argparse.Namespace) -> None:
    ensure_dirs()
    slug = slugify(args.slug)
    data = load_record(slug)
    data["child"] = {
        "sessionKey": args.child_session_key,
        "runId": args.run_id,
        "requestedModel": args.requested_model,
        "modelApplied": args.model_applied,
        "claimedAt": now_iso(),
    }
    data["status"] = "child_claimed"
    save_record(slug, data)
    print_json({"ok": True, "slug": slug, "status": data["status"], "child": data["child"]})


def cmd_capture_draft(args: argparse.Namespace) -> None:
    ensure_dirs()
    slug = slugify(args.slug)
    data = load_record(slug)
    src = Path(args.draft_file)
    if not src.exists():
        raise SystemExit(f"draft file not found: {src}")

    ext = ".json" if args.format == "json" else ".txt"
    dest = DRAFTS / f"{slug}{ext}"
    shutil.copyfile(src, dest)

    data["draft"] = {
        "source": args.source,
        "model": args.model,
        "format": args.format,
        "path": str(dest),
        "bytes": dest.stat().st_size,
        "sha256": sha256_file(dest),
        "capturedAt": now_iso(),
    }
    data["status"] = "draft_captured"
    save_record(slug, data)
    print_json({"ok": True, "slug": slug, "status": data["status"], "draft": data["draft"]})


def cmd_show(args: argparse.Namespace) -> None:
    slug = slugify(args.slug)
    print_json(load_record(slug))


def cmd_verify(args: argparse.Namespace) -> None:
    slug = slugify(args.slug)
    data = load_record(slug)
    errors: list[str] = []

    require_model = args.require_model.strip() if args.require_model else data.get("requireModel", "gpt-mini")

    child = data.get("child") or {}
    draft = data.get("draft") or {}

    if not child:
        errors.append("missing child receipt")
    else:
        if not child.get("modelApplied"):
            errors.append("child receipt does not confirm modelApplied=true")
        child_model = child.get("requestedModel")
        if require_model and child_model != require_model:
            errors.append(f"child requestedModel mismatch: expected {require_model}, got {child_model}")
        if require_model in MINI_ALIASES and not is_mini(child_model):
            errors.append(f"child requestedModel is not mini: {child_model}")

    if not draft:
        errors.append("missing captured draft")
    else:
        draft_path = Path(draft.get("path", ""))
        if draft.get("source") != "child":
            errors.append(f"draft source must be 'child', got {draft.get('source')}")
        if require_model and draft.get("model") != require_model:
            errors.append(f"draft model mismatch: expected {require_model}, got {draft.get('model')}")
        if require_model in MINI_ALIASES and not is_mini(draft.get("model")):
            errors.append(f"draft model is not mini: {draft.get('model')}")
        if not draft_path.exists():
            errors.append(f"captured draft path missing: {draft_path}")
        elif draft_path.stat().st_size == 0:
            errors.append(f"captured draft is empty: {draft_path}")
        elif draft.get("sha256") != sha256_file(draft_path):
            errors.append("captured draft hash mismatch")

    ok = not errors
    if ok:
        data["status"] = "verified"
        data["verifiedAt"] = now_iso()
        save_record(slug, data)

    print_json(
        {
            "ok": ok,
            "slug": slug,
            "requireModel": require_model,
            "status": data.get("status"),
            "errors": errors,
            "record": str(record_path(slug)),
            "draftPath": draft.get("path"),
        },
        ok_exit=ok,
    )


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description="Guardrail for mini-only drafting provenance")
    sub = p.add_subparsers(dest="cmd", required=True)

    p_init = sub.add_parser("init", help="Create a new mini-only task record")
    p_init.add_argument("--slug", required=True)
    p_init.add_argument("--task-kind", required=True)
    p_init.add_argument("--requested-by", required=True)
    p_init.add_argument("--require-model", default="gpt-mini")
    p_init.add_argument("--note")
    p_init.add_argument("--force", action="store_true")
    p_init.set_defaults(func=cmd_init)

    p_child = sub.add_parser("claim-child", help="Store accepted child-run receipt")
    p_child.add_argument("--slug", required=True)
    p_child.add_argument("--child-session-key", required=True)
    p_child.add_argument("--run-id", required=True)
    p_child.add_argument("--requested-model", required=True)
    p_child.add_argument("--model-applied", action="store_true")
    p_child.set_defaults(func=cmd_claim_child)

    p_capture = sub.add_parser("capture-draft", help="Copy exact child output into managed storage")
    p_capture.add_argument("--slug", required=True)
    p_capture.add_argument("--draft-file", required=True)
    p_capture.add_argument("--source", choices=["child", "parent", "manual"], required=True)
    p_capture.add_argument("--model", required=True)
    p_capture.add_argument("--format", choices=["text", "json"], default="text")
    p_capture.set_defaults(func=cmd_capture_draft)

    p_verify = sub.add_parser("verify", help="Verify the mini-only receipt before send")
    p_verify.add_argument("--slug", required=True)
    p_verify.add_argument("--require-model")
    p_verify.set_defaults(func=cmd_verify)

    p_show = sub.add_parser("show", help="Print the stored record")
    p_show.add_argument("--slug", required=True)
    p_show.set_defaults(func=cmd_show)

    return p


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()