#!/usr/bin/env python3
"""Root-owned, tightly allowlisted Academy operator control helper."""
from __future__ import annotations

import hashlib
import json
import os
import pwd
import socket
import subprocess
import sys
import tempfile
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

EXPECTED_HOST = "Work"
EXPECTED_SENTINEL_SHA = "4fa53399f013e76ab05b3a4f093a1a7a4d5b48279a5e7a2a336a159a9c271d11"

REPO = Path("/home/shire3d/ARMOR")
PYTHON = REPO / "venv/bin/python"
RUNUSER = Path("/usr/sbin/runuser")
SCHEDULER = Path("/home/ray/shire-academy-safe-autonomous")
SENTINEL = REPO / "modules/sentinel.py"
RUNTIME = REPO / "data/academy/design_mastery/runtime"
AUTOPILOT = RUNTIME / "autopilot_state.json"
PRACTICE = RUNTIME / "practice_state.json"
SPRINT = RUNTIME / "commercial_design_sprint.json"
INDEX = RUNTIME / "academy_index.sqlite3"
CONTROL = Path("/home/ray/.local/state/shire-unified-learning/control.json")
PLANS_ROOT = Path("/SHiREVault/SHiREAcademy/UnifiedLearning/Plans")
EVIDENCE_ROOT = Path("/SHiREVault/SHiREAcademy/UnifiedLearning/ControlEvidence")
BACKUP_ROOT = Path("/SHiREVault/Backup/OSBackups")

TIMER = "shire-academy-safe-autonomous.timer"
SCHEDULER_SERVICE = "shire-academy-safe-autonomous.service"
UNIFIED_SERVICE = "shire-unified-learning.service"
ACADEMY_UI_SERVICE = "shire-academy-ui.service"
OLD_WORKER = "shire-academy-autonomous-practice.service"

ALLOWED_ACTIONS = {"status", "activate", "pause", "resume", "emergency-stop"}
ALLOWED_AUTOPILOT_CHANGES = {
    "enabled", "paused", "review_required", "safety_stop",
    "safety_stop_reason", "updated_at",
}


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


def fail(message: str, code: int = 1) -> None:
    print(json.dumps({"ok": False, "error": message}, sort_keys=True))
    raise SystemExit(code)


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def load_json(path: Path, default: Any = None) -> Any:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except FileNotFoundError:
        return default


def atomic_json(path: Path, payload: Any, owner: str | None = None, mode: int = 0o600) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, raw = tempfile.mkstemp(prefix=f".{path.name}.", dir=str(path.parent))
    temp = Path(raw)
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as handle:
            json.dump(payload, handle, indent=2, sort_keys=True)
            handle.write("\n")
            handle.flush()
            os.fsync(handle.fileno())
        os.chmod(temp, mode)
        if owner:
            account = pwd.getpwnam(owner)
            try:
                os.chown(temp, account.pw_uid, account.pw_gid)
            except OSError:
                # CIFS mounts may enforce a fixed uid/gid and reject chown even
                # though the resulting file is writable by the service user.
                stat = temp.stat()
                if stat.st_uid not in {account.pw_uid, 0}:
                    raise
        os.replace(temp, path)
    finally:
        try:
            temp.unlink()
        except FileNotFoundError:
            pass


def run(command: list[str], *, check: bool = True, timeout: int = 30) -> subprocess.CompletedProcess[str]:
    result = subprocess.run(command, capture_output=True, text=True, timeout=timeout, check=False)
    if check and result.returncode != 0:
        detail = (result.stderr or result.stdout or "command failed").strip()
        fail(f"{command[0]} failed: {detail}")
    return result


def service_state(name: str) -> str:
    result = run(["systemctl", "is-active", name], check=False, timeout=8)
    return (result.stdout or "").strip() or "inactive"


def service_enabled(name: str) -> str:
    result = run(["systemctl", "is-enabled", name], check=False, timeout=8)
    return (result.stdout or "").strip() or "disabled"


def default_control() -> dict[str, Any]:
    return {
        "schema": "shire.unified_learning.control.v1",
        "version": "SHIRE-Unified-Learning-0005",
        "active": True,
        "paused": False,
        "emergency_stop": False,
        "reason": "",
        "updated_at": utc_now(),
    }


def read_control() -> dict[str, Any]:
    state = load_json(CONTROL, {}) or {}
    if not isinstance(state, dict):
        state = {}
    merged = default_control()
    merged.update(state)
    merged["active"] = bool(merged.get("active"))
    merged["paused"] = bool(merged.get("paused"))
    merged["emergency_stop"] = bool(merged.get("emergency_stop"))
    merged["reason"] = str(merged.get("reason") or "")
    return merged


def write_control(*, active: bool, paused: bool, emergency: bool, reason: str) -> dict[str, Any]:
    state = default_control()
    state.update({
        "active": active,
        "paused": paused,
        "emergency_stop": emergency,
        "reason": reason,
        "updated_at": utc_now(),
    })
    atomic_json(CONTROL, state, owner="ray")
    return state


def scheduler(action: str) -> None:
    # The helper needs root only for systemd. Run the user-owned scheduler as
    # ray so pause/resume state files never become root-owned.
    run([
        str(RUNUSER),
        "-u",
        "ray",
        "--",
        "/usr/bin/env",
        "HOME=/home/ray",
        "PATH=/home/ray:/home/ray/.local/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
        str(PYTHON),
        str(SCHEDULER),
        action,
    ], timeout=30)


def verify_environment() -> None:
    if os.geteuid() != 0:
        fail("Academy control helper must run through the installed sudo rule.")
    if socket.gethostname() != EXPECTED_HOST:
        fail(f"Expected host {EXPECTED_HOST}.")
    for path in (PYTHON, RUNUSER, SCHEDULER, SENTINEL, AUTOPILOT, PRACTICE, SPRINT, INDEX):
        if not path.exists():
            fail(f"Required path is missing: {path}")
    if sha256(SENTINEL) != EXPECTED_SENTINEL_SHA:
        fail("Sentinel source hash does not match the protected baseline.")
    mount = run(
        ["findmnt", "-T", str(BACKUP_ROOT), "-rn", "-o", "FSTYPE,SOURCE"],
        check=False,
        timeout=8,
    )
    rows = [line.split(None, 1) for line in mount.stdout.splitlines() if line.strip()]
    if not any(row and row[0] == "cifs" for row in rows):
        fail("SHiREVault is not on the required CIFS network mount.")
    if service_state(OLD_WORKER) == "active":
        fail("Retired Academy worker is unexpectedly active.")


def protected_snapshot() -> dict[str, str]:
    return {str(path): sha256(path) for path in (PRACTICE, SPRINT, INDEX)}


def autopilot_without_allowed(path: Path) -> dict[str, Any]:
    value = load_json(path, {}) or {}
    if not isinstance(value, dict):
        fail(f"Invalid autopilot state: {path}")
    return {key: child for key, child in value.items() if key not in ALLOWED_AUTOPILOT_CHANGES}


def mark_active_jobs_safety_stopped() -> int:
    count = 0
    if not PLANS_ROOT.exists():
        return count
    active = {"queued", "running", "paused_by_operator"}
    for job_path in PLANS_ROOT.glob("ULP-*/jobs/ULJ-*/job.json"):
        job = load_json(job_path, {}) or {}
        if not isinstance(job, dict) or job.get("status") not in active:
            continue
        stopped_at = utc_now()
        failure = {
            "phase": "operator_emergency_stop",
            "failed_at": stopped_at,
            "error": "Emergency Stop was activated from the ARMOR OS Academy dashboard.",
            "preserved": True,
            "operator_emergency_stop": True,
        }
        failures = job_path.parent / "failures"
        failures.mkdir(parents=True, exist_ok=True)
        atomic_json(
            failures / f"{datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ')}-{uuid.uuid4().hex[:8]}.json",
            failure,
            owner="ray",
            mode=0o640,
        )
        job.update({
            "status": "safety_stop",
            "failure": failure,
            "updated_at": stopped_at,
            "operator_control": {
                "emergency_stop": True,
                "stopped_at": stopped_at,
            },
        })
        atomic_json(job_path, job, owner="ray", mode=0o640)
        count += 1
    return count


def write_action_receipt(
    action: str,
    before: dict[str, Any],
    after: dict[str, Any],
    stopped_jobs: int,
) -> str:
    day = datetime.now(timezone.utc).strftime("%Y-%m-%d")
    stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    path = EVIDENCE_ROOT / day / f"{stamp}-{action}-{uuid.uuid4().hex[:8]}.json"
    payload = {
        "schema": "shire.academy.operator_control_receipt.v1",
        "version": "SHIRE-Unified-Learning-0005",
        "recorded_at": utc_now(),
        "host": socket.gethostname(),
        "operator": "ray",
        "action": action,
        "control_before": before,
        "control_after": after,
        "stopped_jobs": stopped_jobs,
        "timer_state": service_state(TIMER),
        "unified_service_state": service_state(UNIFIED_SERVICE),
        "safe_scheduler_service_state": service_state(SCHEDULER_SERVICE),
        "retired_worker_state": service_state(OLD_WORKER),
        "sentinel_sha256": sha256(SENTINEL),
        "sentinel_unchanged": True,
        "protected_academy_state_changed": False,
        "forge_unlock_allowed": False,
        "counted_attempts_changed": False,
        "certification_changed": False,
    }
    atomic_json(path, payload, owner="ray", mode=0o640)
    return str(path)


def result_payload(
    action: str,
    state: dict[str, Any],
    stopped_jobs: int = 0,
    receipt: str = "",
) -> dict[str, Any]:
    return {
        "ok": True,
        "action": action,
        "control": state,
        "stopped_jobs": stopped_jobs,
        "receipt": receipt,
        "timer": {
            "active": service_state(TIMER),
            "enabled": service_enabled(TIMER),
        },
        "services": {
            "unified_learning": service_state(UNIFIED_SERVICE),
            "academy_ui": service_state(ACADEMY_UI_SERVICE),
            "safe_scheduler": service_state(SCHEDULER_SERVICE),
            "retired_worker": service_state(OLD_WORKER),
        },
        "forge_unlock_allowed": False,
        "protected_academy_state_changed": False,
        "sentinel_unchanged": True,
    }


def main() -> None:
    if len(sys.argv) != 2 or sys.argv[1] not in ALLOWED_ACTIONS:
        fail("Allowed actions: status, activate, pause, resume, emergency-stop.", 2)
    action = sys.argv[1]
    verify_environment()

    if action == "status":
        print(json.dumps(result_payload(action, read_control()), sort_keys=True))
        return

    control_before = read_control()
    protected_before = protected_snapshot()
    sentinel_before = sha256(SENTINEL)
    autopilot_before = autopilot_without_allowed(AUTOPILOT)
    stopped_jobs = 0

    if action == "activate":
        write_control(active=True, paused=False, emergency=False, reason="")
        run(["systemctl", "start", UNIFIED_SERVICE])
        run(["systemctl", "start", ACADEMY_UI_SERVICE])
        scheduler("resume")
        run(["systemctl", "enable", "--now", TIMER])

    elif action == "pause":
        state = read_control()
        if state["emergency_stop"] or not state["active"]:
            fail("Academy is inactive. Use Activate Academy first.")
        write_control(
            active=True,
            paused=True,
            emergency=False,
            reason="Paused by the operator from ARMOR OS Academy.",
        )
        scheduler("pause")
        run(["systemctl", "stop", TIMER])

    elif action == "resume":
        state = read_control()
        if state["emergency_stop"] or not state["active"]:
            fail("Emergency Stop is active. Use Activate Academy instead of Resume.")
        run(["systemctl", "start", UNIFIED_SERVICE])
        scheduler("resume")
        write_control(active=True, paused=False, emergency=False, reason="")
        run(["systemctl", "start", TIMER])

    elif action == "emergency-stop":
        write_control(
            active=False,
            paused=True,
            emergency=True,
            reason="Emergency Stop activated by the operator from ARMOR OS Academy.",
        )
        scheduler("pause")
        run(["systemctl", "stop", TIMER], check=False)
        run(["systemctl", "stop", SCHEDULER_SERVICE], check=False)
        run(["systemctl", "stop", UNIFIED_SERVICE], check=False)
        stopped_jobs = mark_active_jobs_safety_stopped()
        # Restart only the read/control service. Execution stays blocked by control.json.
        run(["systemctl", "start", UNIFIED_SERVICE])

    if protected_snapshot() != protected_before:
        fail("Protected practice, sprint or Academy index state changed during the control action.")
    if autopilot_without_allowed(AUTOPILOT) != autopilot_before:
        fail("Academy control changed autopilot fields outside the approved pause controls.")
    if sha256(SENTINEL) != sentinel_before or sentinel_before != EXPECTED_SENTINEL_SHA:
        fail("Sentinel changed during the Academy control action.")
    if service_state(OLD_WORKER) == "active":
        fail("Retired Academy worker became active.")

    state = read_control()
    if action in {"activate", "resume"}:
        if state["paused"] or state["emergency_stop"] or not state["active"]:
            fail("Academy did not enter the active state.")
        if service_state(TIMER) != "active":
            fail("Safe Academy timer did not become active.")
    elif action == "pause":
        if not state["paused"] or state["emergency_stop"] or not state["active"]:
            fail("Academy did not enter the paused state.")
        if service_state(TIMER) == "active":
            fail("Safe Academy timer remained active after Pause.")
    elif action == "emergency-stop":
        if not state["emergency_stop"] or state["active"]:
            fail("Emergency Stop state was not preserved.")
        if service_state(TIMER) == "active" or service_state(SCHEDULER_SERVICE) == "active":
            fail("Academy execution remained active after Emergency Stop.")
        if service_state(UNIFIED_SERVICE) != "active":
            fail("Unified Learning control/dashboard service did not recover after Emergency Stop.")

    receipt = write_action_receipt(action, control_before, state, stopped_jobs)
    print(json.dumps(result_payload(action, state, stopped_jobs, receipt), sort_keys=True))


if __name__ == "__main__":
    main()
