#!/usr/bin/env python3
"""AI Teachers Agentic OS Lab: a safe, local-only control-plane simulation.

Python 3.10+ and standard library only. No model calls, sockets, shell commands,
file mutations outside the selected state directory, or external side effects.
Use invented training notes only. This is NOT a secure production sandbox.

The four functions are illustrative stages, not OS-level isolation.
True planner/worker/verifier separation needs separate runtime identities,
permissions, durable immutable storage, and real independent checks.
"""
from __future__ import annotations

import argparse
import datetime as dt
import hashlib
import json
import re
import sys
import uuid
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parent
DEFAULT_POLICY = ROOT / "constitution.json"
SAFE_ID = re.compile(r"^[a-zA-Z0-9_-]{1,70}$")
UNTRUSTED_COMMAND = re.compile(
    r"(ignore (all )?(prior|previous) instructions|"
    r"mark (this |the )?(task |output )?(as )?(passed|approved|verified)|"
    r"(system|developer) override|"
    r"reveal (the )?(secret|api key)|"
    r"send (all )?(files|passwords) to)",
    re.IGNORECASE,
)


class PolicyError(Exception):
    """Deny by default when a policy or input cannot be validated."""


def utc_now() -> str:
    return dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds")


def load_json(path: Path) -> dict[str, Any]:
    value = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(value, dict):
        raise PolicyError(f"Expected object in {path.name}")
    return value


def canonical(data: Any) -> str:
    return json.dumps(data, sort_keys=True, ensure_ascii=False, separators=(",", ":"))


def validate_policy(policy: dict[str, Any]) -> None:
    required = ("daily_budget_usd", "max_estimated_task_usd",
                "max_input_tokens", "max_output_tokens",
                "min_audited_runs", "min_audited_pass_rate")
    for key in required:
        if isinstance(policy.get(key), bool) or not isinstance(policy.get(key), (int, float)):
            raise PolicyError("Invalid policy field: " + key)
    if not (policy["daily_budget_usd"] > 0
            and 0 < policy["max_estimated_task_usd"] <= policy["daily_budget_usd"]
            and policy["max_input_tokens"] > 0
            and policy["max_output_tokens"] > 0
            and policy["min_audited_runs"] >= 1
            and 0 < policy["min_audited_pass_rate"] <= 1):
        raise PolicyError("Out-of-range policy fields")
    if policy.get("allow_external_actions") is not False:
        raise PolicyError("The teaching runtime requires external actions disabled")
    if not isinstance(policy.get("allowed_task_types"), list) or not policy["allowed_task_types"]:
        raise PolicyError("allowed_task_types must be a nonempty list")
    rates = policy.get("illustrative_token_usd_per_million")
    if not isinstance(rates, dict) or not all(
        isinstance(rates.get(k), (int, float)) and rates[k] > 0
        for k in ("input", "output")
    ):
        raise PolicyError("Illustrative token rates are missing")


def validate_task(task: dict[str, Any], policy: dict[str, Any]) -> None:
    if not isinstance(task.get("id"), str) or not SAFE_ID.fullmatch(task["id"]):
        raise PolicyError("Task id must contain letters, numbers, '-' or '_'")
    if task.get("type") not in policy["allowed_task_types"]:
        raise PolicyError("Task type not approved by the policy")
    records = task.get("records")
    if not isinstance(records, list) or not 1 <= len(records) <= 15:
        raise PolicyError("Use 1–15 fictional, structured source records")
    ids: set[str] = set()
    for item in records:
        if not isinstance(item, dict) or not isinstance(item.get("id"), str):
            raise PolicyError("Every record needs an id")
        if not SAFE_ID.fullmatch(item["id"]) or item["id"] in ids:
            raise PolicyError("Source ids must be unique and safe")
        ids.add(item["id"])
        if not isinstance(item.get("text"), str) or not 1 <= len(item["text"]) <= 500:
            raise PolicyError("Each source needs 1–500 characters of text")
    for field, limit in (("estimated_input_tokens", policy["max_input_tokens"]),
                         ("estimated_output_tokens", policy["max_output_tokens"])):
        count = task.get(field)
        if isinstance(count, bool) or not isinstance(count, int) or not 1 <= count <= limit:
            raise PolicyError(f"Invalid or excessive {field}")


def estimated_cost(task: dict[str, Any], policy: dict[str, Any]) -> float:
    """Illustrative reservation, NOT actual measured token usage or vendor spend."""
    rates = policy["illustrative_token_usd_per_million"]
    return round((task["estimated_input_tokens"] * rates["input"] +
                  task["estimated_output_tokens"] * rates["output"]) / 1_000_000, 8)


def events(state: Path) -> list[dict[str, Any]]:
    path = state / "events.jsonl"
    if not path.exists():
        return []
    out: list[dict[str, Any]] = []
    for line in path.read_text(encoding="utf-8").splitlines():
        if not line.strip():
            continue
        entry = json.loads(line)
        if not isinstance(entry, dict) or entry.get("kind") not in ("run", "audit", "monitor_alert"):
            raise PolicyError("Unrecognized ledger event: fail closed")
        out.append(entry)
    return out


def append_event(state: Path, value: dict[str, Any]) -> None:
    state.mkdir(parents=True, exist_ok=True)
    with (state / "events.jsonl").open("a", encoding="utf-8") as file:
        file.write(canonical(value) + "\n")


def daily_spend(items: list[dict[str, Any]], date: str) -> float:
    return round(sum(
        item.get("estimated_cost_usd", 0)
        for item in items if item["kind"] == "run"
        and item.get("charged") is True
        and item["timestamp"].startswith(date)
    ), 8)


def plan_read_only(task: dict[str, Any]) -> dict[str, Any]:
    """A constrained, pure function. It cannot write files or call tools."""
    return {
        "task_id": task["id"],
        "task_type": task["type"],
        "source_ids": [item["id"] for item in task["records"]],
        "allowed_output": "LOCAL_DRAFT_ONLY",
        "human_approval": "REQUIRED",
    }


def execute_draft(task: dict[str, Any], plan: dict[str, Any]) -> dict[str, Any]:
    """Simulated worker: no LLM, no network, no shell, no external writes."""
    claims = [
        {"source_id": item["id"], "quote": item["text"]}
        for item in task["records"]
    ]
    return {
        "task_id": plan["task_id"],
        "label": "SIMULATED DRAFT — NOT APPROVED",
        "summary": f"Local demonstration using {len(claims)} fictional records.",
        "claims": claims,
        "external_actions": [],
    }


def verify_independently(task: dict[str, Any], draft: dict[str, Any]) -> dict[str, Any]:
    """Deterministic teaching verifier; no trust in the worker's own assertions."""
    reasons: list[str] = []
    source = {item["id"]: item["text"] for item in task["records"]}
    claims = draft.get("claims")
    if not isinstance(claims, list) or len(claims) != len(source):
        reasons.append("Missing or extra source claims")
        claims = []
    cited: set[str] = set()
    for claim in claims:
        if not isinstance(claim, dict):
            reasons.append("Claim not structured")
            continue
        sid, quote = claim.get("source_id"), claim.get("quote")
        if sid in cited:
            reasons.append("Duplicate source id")
        cited.add(sid)
        if sid not in source or quote != source.get(sid):
            reasons.append("Claim does not match original source")
        if not isinstance(quote, str) or UNTRUSTED_COMMAND.search(quote):
            reasons.append("Suspicious command inside untrusted source material")
    if cited != set(source):
        reasons.append("Evidence coverage is incomplete")
    if draft.get("external_actions") != []:
        reasons.append("Forbidden external actions requested")
    if draft.get("label") != "SIMULATED DRAFT — NOT APPROVED":
        reasons.append("Output pretends to have an approval")
    return {"passed": not reasons, "reasons": sorted(set(reasons)),
            "source_count": len(source), "claim_count": len(claims)}


def trust(items: list[dict[str, Any]], task_type: str, policy: dict[str, Any]) -> dict[str, Any]:
    audited = [
        e for e in items if e["kind"] == "audit"
        and e["task_type"] == task_type
    ]
    failures = [e for e in audited if e["result"] == "fail"]
    alerts = [e for e in items if e["kind"] == "monitor_alert"
              and e["task_type"] == task_type]
    passed = sum(e["result"] == "pass" for e in audited)
    ratio = round(passed / len(audited), 3) if audited else 0.0
    # A failed human review or monitor alert demands explicit reconciliation.
    # No demo command can silently reset this marker.
    if failures or alerts:
        state = "REVOKED_REQUIRES_REVIEW"
    elif (len(audited) >= policy["min_audited_runs"]
          and ratio >= policy["min_audited_pass_rate"]):
        state = "ELIGIBLE_FOR_HUMAN_REVIEW_OF_NARROW_AUTONOMY"
    else:
        state = "DRAFT_ONLY"
    return {"task_type": task_type, "audited_runs": len(audited),
            "passed_audits": passed, "pass_rate": ratio, "state": state,
            "external_actions_allowed": False}


def run(task_file: Path, state: Path, policy: dict[str, Any]) -> dict[str, Any]:
    task = load_json(task_file)
    validate_task(task, policy)
    log = events(state)
    fp = hashlib.sha256(canonical(task).encode()).hexdigest()
    if any(e["kind"] == "run" and e.get("task_id") == task["id"] for e in log):
        raise PolicyError("Duplicate task id; use a new id for each deliberate run")
    cost = estimated_cost(task, policy)
    day = utc_now()[:10]
    spent = daily_spend(log, day)
    if (cost > policy["max_estimated_task_usd"]
            or spent + cost > policy["daily_budget_usd"]):
        event = {"kind": "run", "timestamp": utc_now(), "run_id": str(uuid.uuid4()),
                 "task_id": task["id"], "task_type": task["type"],
                 "status": "BLOCKED_BUDGET", "charged": False,
                 "estimated_cost_usd": 0.0, "input_hash": fp,
                 "human_approval": False}
        append_event(state, event)
        return {"status": "BLOCKED_BUDGET", "estimated_cost_usd": cost,
                "remaining_budget_usd": round(policy["daily_budget_usd"] - spent, 8)}
    plan = plan_read_only(task)
    draft = execute_draft(task, plan)
    verification = verify_independently(task, draft)
    event = {"kind": "run", "timestamp": utc_now(), "run_id": str(uuid.uuid4()),
             "task_id": task["id"], "task_type": task["type"],
             "status": "VERIFIED_DRAFT_NEEDS_HUMAN_REVIEW" if verification["passed"] else "REJECTED",
             "charged": True, "estimated_cost_usd": cost, "input_hash": fp,
             "policy_version": policy.get("version"), "source_snapshot": task,
             "plan": plan, "draft": draft, "verification": verification,
             "human_approval": False}
    append_event(state, event)
    return {"task_id": task["id"], "status": event["status"],
            "estimated_cost_usd": cost, "verification": verification,
            "human_approval": False, "external_actions": []}


def audit(task_id: str, reviewer: str, result: str,
          state: Path, policy: dict[str, Any]) -> dict[str, Any]:
    if not SAFE_ID.fullmatch(task_id) or not SAFE_ID.fullmatch(reviewer):
        raise PolicyError("Invalid id or reviewer label")
    items = events(state)
    matches = [e for e in items if e["kind"] == "run" and e["task_id"] == task_id]
    if len(matches) != 1:
        raise PolicyError("Audits require exactly one existing run")
    record = matches[0]
    if record["status"] not in ("VERIFIED_DRAFT_NEEDS_HUMAN_REVIEW", "REJECTED"):
        raise PolicyError("Cannot audit a budget-blocked task")
    if any(e["kind"] == "audit" and e["task_id"] == task_id for e in items):
        raise PolicyError("Duplicate audit is forbidden")
    if result == "pass" and record["status"] != "VERIFIED_DRAFT_NEEDS_HUMAN_REVIEW":
        raise PolicyError("Cannot pass a draft the verifier rejected")
    append_event(state, {"kind": "audit", "timestamp": utc_now(),
                         "task_id": task_id, "task_type": record["task_type"],
                         "reviewer": reviewer, "result": result})
    return trust(events(state), record["task_type"], policy)


def monitor(state: Path, policy: dict[str, Any]) -> list[dict[str, Any]]:
    log = events(state)
    rows: list[dict[str, Any]] = []
    for entry in [e for e in log if e["kind"] == "run" and e.get("source_snapshot")]:
        check = verify_independently(entry["source_snapshot"], entry["draft"])
        result = {"task_id": entry["task_id"], "status": "CHECKED" if check["passed"]
                  else "ALERT_REVIEW_REQUIRED", "verification": check}
        rows.append(result)
        if not check["passed"] and not any(
            e["kind"] == "monitor_alert" and e.get("task_id") == entry["task_id"] for e in log
        ):
            append_event(state, {"kind": "monitor_alert", "timestamp": utc_now(),
                                 "task_id": entry["task_id"], "task_type": entry["task_type"],
                                 "reason": "Saved evidence no longer verifies"})
    return rows


def report(state: Path, policy: dict[str, Any]) -> dict[str, Any]:
    log = events(state)
    today = utc_now()[:10]
    return {"date_utc": today, "simulated_daily_estimated_spend_usd": daily_spend(log, today),
            "simulated_daily_cap_usd": policy["daily_budget_usd"],
            "run_count": sum(e["kind"] == "run" for e in log),
            "human_audit_count": sum(e["kind"] == "audit" for e in log),
            "trust": [trust(log, task_type, policy)
                      for task_type in policy["allowed_task_types"]],
            "disclaimer": "Simulation only: not actual usage, immutable audit or an API spend cap."}


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--policy", type=Path, default=DEFAULT_POLICY)
    parser.add_argument("--state", type=Path, default=Path("./agentic-os-state"))
    commands = parser.add_subparsers(dest="command", required=True)
    run_parser = commands.add_parser("run", help="Simulate planner, worker, independent verifier and gate")
    run_parser.add_argument("--task", required=True, type=Path)
    audit_parser = commands.add_parser("audit", help="Manually record a human review")
    audit_parser.add_argument("--task-id", required=True)
    audit_parser.add_argument("--reviewer", required=True)
    audit_parser.add_argument("--result", required=True, choices=("pass", "fail"))
    commands.add_parser("monitor", help="Recheck stored training evidence, without network access")
    commands.add_parser("report", help="Show local status, estimated budget and earned-trust records")
    args = parser.parse_args()
    try:
        policy = load_json(args.policy)
        validate_policy(policy)
        if args.command == "run":
            output = run(args.task, args.state, policy)
        elif args.command == "audit":
            output = audit(args.task_id, args.reviewer, args.result, args.state, policy)
        elif args.command == "monitor":
            output = monitor(args.state, policy)
        else:
            output = report(args.state, policy)
    except (PolicyError, OSError, ValueError, json.JSONDecodeError) as exc:
        print("DENIED:", str(exc), file=sys.stderr)
        return 2
    print(json.dumps(output, indent=2, ensure_ascii=False))
    return 0


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