#!/usr/bin/env python3
"""Durable local session/job state paths, validation, and publication."""

from __future__ import annotations

import contextlib
import copy
import json
import os
import re
from collections.abc import Mapping
from pathlib import Path
from typing import Any

from mmo_util import (
    SAFE_JOB_ID,
    append_jsonl,
    atomic_write_json,
    is_within,
    process_alive,
    process_group_alive,
    process_matches,
    read_json_object,
    state_root,
    terminate_process,
    terminate_process_group,
    utc_now,
)
from mmo_version import MMO_SCHEMA_VERSION, PACKAGE_VERSION

ACTIVE_JOB_STATUSES = {
    "queued",
    "starting",
    "running",
    "waiting",
    "paused",
    "detached",
    "recovering",
    "finalizing",
    "cancelling",
}
RECOVERABLE_JOB_STATUSES = ACTIVE_JOB_STATUSES | {"suspended"}
ADMITTING_JOB_STATUSES = {"queued", "starting", "running", "waiting", "paused", "detached"}
TERMINAL_JOB_STATUSES = {
    "completed",
    "completed_with_warnings",
    "failed",
    "stopped",
    "cancelled",
}
ADMITTING_SESSION_STATUSES = {"starting", "running", "detached", "paused", "suspended"}
ACTIVE_SESSION_STATUSES = {
    "starting",
    "running",
    "detached",
    "paused",
    "suspended",
    "finishing",
    "stopping",
    "cancelling",
}
TERMINAL_SESSION_STATUSES = {"completed", "stopped", "failed", "cancelled"}

ROOT_EXECUTION_HOSTS = frozenset({"app_server"})
RETIRED_SESSION_FIELDS = frozenset(
    {
        "root_execution_policy_enforced",
        "root_rollout_path",
    }
)
RETIRED_JOB_FIELDS = frozenset(
    {
        "execution_policy",
        "renewal_quantum_seconds",
        "max_active_work_seconds",
        "active_work_cap_seconds",
        "active_work_seconds",
        "automatic_renewal_enabled",
        "app_server_rollout_path",
    }
)


def _validate_switchyard_identity(data: Mapping[str, Any], label: str) -> None:
    if "switchyard_version" not in data:
        raise ValueError(f"{label} has no Switchyard version identity")
    version = data.get("switchyard_version")
    if data.get("gateway_base_url") is None:
        if version is not None:
            raise ValueError(f"{label} has a Switchyard version without a gateway")
        return
    if not isinstance(version, str) or re.fullmatch(r"[0-9]+\.[0-9]+\.[0-9]+", version) is None:
        raise ValueError(f"{label} Switchyard version is invalid")


def root_mcp_token(session_id: str) -> str:
    """Read the canonical persisted root MCP capability for an active session."""

    try:
        data = read_json_object(
            session_dir(session_id) / "capabilities.json",
            label="session capabilities",
        )
    except FileNotFoundError as exc:
        raise RuntimeError("root MCP caller capability is unavailable for this session") from exc
    token = data.get("root")
    if not isinstance(token, str) or not token:
        raise RuntimeError("root MCP caller capability is invalid")
    return token


def store_session_capabilities(
    directory: Path,
    *,
    root_token: str,
    native_tokens: Mapping[str, str],
) -> None:
    """Persist host-scoped MCP identities needed by detached app-server processes."""

    if not root_token or not all(
        isinstance(value, str) and value for value in native_tokens.values()
    ):
        raise ValueError("session capabilities must be non-empty strings")
    atomic_write_json(
        directory / "capabilities.json",
        {
            "schema_version": MMO_SCHEMA_VERSION,
            "root": root_token,
            "native": dict(native_tokens),
        },
        0o600,
    )


def load_session_capabilities(directory: Path) -> tuple[str, dict[str, str]]:
    data = read_json_object(directory / "capabilities.json", label="session capabilities")
    root_token = data.get("root")
    native_tokens = data.get("native")
    if (
        data.get("schema_version") != MMO_SCHEMA_VERSION
        or not isinstance(root_token, str)
        or not root_token
        or not isinstance(native_tokens, Mapping)
        or not all(
            isinstance(key, str) and isinstance(value, str) and value
            for key, value in native_tokens.items()
        )
    ):
        raise ValueError("session capabilities are invalid")
    return root_token, dict(native_tokens)


def root_mcp_capability_environment(session: Mapping[str, Any]) -> dict[str, str]:
    """Return the authenticated Agent-MCP identity for a root-owned process."""

    session_id = str(session["session_id"])
    return {
        "MMO_ROOT_SESSION_ID": session_id,
        "MMO_RUN_ID": str(session["current_run_id"]),
        "MMO_CALLER_AGENT": str(session["root_agent"]),
        "MMO_CALLER_TOKEN": root_mcp_token(session_id),
    }


def revoke_session_capabilities(session_id: str) -> None:
    """Revoke the bearer identities for a terminal session."""

    # A terminal session retains its immutable transcript and evidence, but its
    # bearer credentials no longer authorize any control-plane operation.
    (session_dir(session_id) / "capabilities.json").unlink(missing_ok=True)


def sessions_root() -> Path:
    root = state_root() / "sessions"
    root.mkdir(parents=True, exist_ok=True, mode=0o700)
    return root


def jobs_root() -> Path:
    root = state_root() / "jobs"
    root.mkdir(parents=True, exist_ok=True, mode=0o700)
    return root


def runtime_lock_path() -> Path:
    return state_root() / ".runtime.lock"


def session_dir(session_id: str) -> Path:
    if not SAFE_JOB_ID.fullmatch(session_id):
        raise ValueError("invalid session id")
    path = sessions_root() / session_id
    if not path.is_dir():
        raise FileNotFoundError(f"unknown session: {session_id}")
    return path


def job_dir(job_id: str) -> Path:
    if not SAFE_JOB_ID.fullmatch(job_id):
        raise ValueError("invalid job id")
    path = jobs_root() / job_id
    if not path.is_dir():
        raise FileNotFoundError(f"unknown agent job: {job_id}")
    return path


def session_state_path(directory: Path) -> Path:
    return directory / "session.json"


def job_state_path(directory: Path) -> Path:
    return directory / "metadata.json"


def ensure_runs_root(directory: Path) -> Path:
    root = directory / "runs"
    if root.is_symlink():
        raise RuntimeError("session runs directory cannot be a symlink")
    root.mkdir(parents=True, exist_ok=True, mode=0o700)
    if not is_within(root.resolve(), directory.resolve()):
        raise RuntimeError("session runs directory escapes its logical session")
    return root


def _run_path(directory: Path, run_id: str) -> Path:
    if not SAFE_JOB_ID.fullmatch(run_id):
        raise ValueError("invalid run id")
    root = directory / "runs"
    if root.is_symlink():
        raise RuntimeError("session runs directory cannot be a symlink")
    path = root / run_id / "run.json"
    if not path.is_file():
        raise FileNotFoundError(f"unknown session run: {run_id}")
    return path


def session_lifecycle_lock_path(directory: Path) -> Path:
    path = directory / "lifecycle.lock"
    if directory.is_symlink() or path.is_symlink():
        raise RuntimeError("persistent session lifecycle lock cannot traverse a symlink")
    if not is_within(path.resolve(), directory.resolve()):
        raise RuntimeError("persistent session lifecycle lock escapes its logical session")
    return path


def job_control_lock_path(directory: Path) -> Path:
    path = directory / "control.lock"
    if directory.is_symlink() or path.is_symlink():
        raise RuntimeError("persistent job control lock cannot traverse a symlink")
    if not is_within(path.resolve(), directory.resolve()):
        raise RuntimeError("persistent job control lock escapes its logical job")
    return path


def _read_run_state(path: Path) -> dict[str, Any]:
    session_directory = path.parent.parent.parent
    if path.parent.parent.is_symlink() or path.parent.is_symlink() or path.is_symlink():
        raise ValueError("session run state cannot traverse a symlink")
    if not is_within(path.resolve(), session_directory.resolve()):
        raise ValueError("session run state escapes its logical session")
    data = read_json_object(path, label="session run state")
    if data.get("schema_version") != MMO_SCHEMA_VERSION or isinstance(
        data.get("schema_version"), bool
    ):
        raise ValueError(f"unsupported session run schema; expected {MMO_SCHEMA_VERSION}")
    if data.get("package_version") != PACKAGE_VERSION:
        raise ValueError(f"persistent session run package must be {PACKAGE_VERSION}")
    sequence = data.get("sequence")
    if not isinstance(sequence, int) or isinstance(sequence, bool) or sequence != 1:
        raise ValueError("immutable session run sequence must be exactly 1")
    if data.get("kind") != "initial":
        raise ValueError("immutable session run kind must be initial")
    if data.get("status") not in ACTIVE_SESSION_STATUSES | TERMINAL_SESSION_STATUSES:
        raise ValueError("session run status is invalid")
    _validate_switchyard_identity(data, "persistent session run")
    expected_id = path.parent.name
    if data.get("run_id") != expected_id:
        raise ValueError(
            "session run identity mismatch: "
            f"directory is {expected_id!r}, record is {data.get('run_id')!r}"
        )
    if data.get("session_id") != path.parent.parent.parent.name:
        raise ValueError("session run belongs to another logical session")
    return data


def load_session_run(session_id: str, run_id: str) -> dict[str, Any]:
    return _read_run_state(_run_path(session_dir(session_id), run_id))


def iter_session_runs(session_id: str) -> list[dict[str, Any]]:
    directory = session_dir(session_id)
    root = directory / "runs"
    if root.is_symlink():
        raise RuntimeError("session runs directory cannot be a symlink")
    if not root.is_dir():
        return []
    results: list[dict[str, Any]] = []
    for candidate in sorted(root.iterdir(), reverse=True):
        path = candidate / "run.json"
        if not candidate.is_dir() or not path.is_file():
            continue
        try:
            results.append(_read_run_state(path))
        except (OSError, ValueError, json.JSONDecodeError) as exc:
            raise RuntimeError(f"invalid session run state: {path}: {exc}") from exc
    return results


def public_run(data: Mapping[str, Any]) -> dict[str, Any]:
    keys = (
        "run_id",
        "package_version",
        "session_id",
        "sequence",
        "kind",
        "status",
        "created_at",
        "started_at",
        "finished_at",
        "root_pid",
        "exit_code",
        "gateway_base_url",
        "switchyard_version",
        "root_execution_host",
        "root_goal_status",
        "root_goal_tokens_used",
        "root_goal_token_budget",
        "error",
    )
    return {key: data.get(key) for key in keys if data.get(key) is not None}


_RUN_MIRROR_FIELDS = (
    "status",
    "started_at",
    "finished_at",
    "root_pid",
    "root_pgid",
    "root_start_token",
    "root_process_group_isolated",
    "exit_code",
    "error",
    "cancel_requested_at",
    "transition_started_at",
    "requested_terminal_status",
    "gateway_base_url",
    "gateway_pid",
    "gateway_hash",
    "gateway_routing_log_path",
    "switchyard_version",
    "route_availability",
    "root_mcp_token_hash",
    "native_token_hashes",
    "root_execution_host",
    "root_execution_mode",
    "root_thread_id",
    "root_thread_generation",
    "root_thread_lineage",
    "root_thread_transition",
    "root_goal_status",
    "root_goal_objective",
    "root_goal_bootstrap_pending",
    "root_goal_token_budget",
    "root_max_goal_token_budget",
    "root_goal_tokens_used",
    "root_goal_time_used_seconds",
    "root_stall_warning_seconds",
    "root_last_progress_at",
    "root_finalization_grace_seconds",
    "root_finalization_started_at",
    "root_app_server_lifecycle_timeout_seconds",
    "root_pending_request_count",
    "root_last_turn_id",
    "root_turn_start_pending",
    "root_finalizing",
    "root_completion_deferred",
    "root_completion_deferred_jobs",
)


def _validate_root_thread_lineage(data: Mapping[str, Any]) -> None:
    generation = data.get("root_thread_generation")
    lineage = data.get("root_thread_lineage")
    transition = data.get("root_thread_transition")
    if not isinstance(generation, int) or isinstance(generation, bool) or generation < 0:
        raise ValueError("persistent session root thread generation is invalid")
    if not isinstance(lineage, list):
        raise ValueError("persistent session root thread lineage must be a list")
    if len(lineage) != generation:
        raise ValueError("persistent session root thread lineage disagrees with its generation")
    seen: set[str] = set()
    for index, raw in enumerate(lineage, 1):
        if not isinstance(raw, Mapping):
            raise ValueError("persistent session root thread lineage entry is invalid")
        thread_id = raw.get("thread_id")
        if (
            raw.get("generation") != index
            or isinstance(raw.get("generation"), bool)
            or not isinstance(thread_id, str)
            or not thread_id
            or thread_id in seen
            or not isinstance(raw.get("codex_session_id"), str)
            or not raw.get("codex_session_id")
            or not isinstance(raw.get("adopted_at"), str)
            or not raw.get("adopted_at")
            or not isinstance(raw.get("reason"), str)
            or not raw.get("reason")
        ):
            raise ValueError("persistent session root thread lineage entry is invalid")
        seen.add(thread_id)
        superseded_at = raw.get("superseded_at")
        successor_id = raw.get("successor_thread_id")
        terminal_entry = index == generation
        if terminal_entry:
            if superseded_at is not None or successor_id is not None:
                raise ValueError("current root thread lineage entry cannot be superseded")
        else:
            successor = lineage[index]
            if (
                not isinstance(successor, Mapping)
                or not isinstance(superseded_at, str)
                or not superseded_at
                or not isinstance(successor_id, str)
                or successor_id != successor.get("thread_id")
            ):
                raise ValueError("superseded root thread lineage entry is incomplete")
    root_thread_id = data.get("root_thread_id")
    if generation == 0:
        if root_thread_id is not None:
            raise ValueError("uninitialized root thread lineage has a current thread")
    elif root_thread_id != lineage[-1].get("thread_id"):
        raise ValueError("persistent session current root thread disagrees with its lineage")
    if transition is not None:
        if not isinstance(transition, Mapping):
            raise ValueError("persistent session root thread transition is invalid")
        if (
            transition.get("from_thread_id") != root_thread_id
            or not isinstance(transition.get("to_thread_id"), str)
            or not transition.get("to_thread_id")
            or transition.get("to_thread_id") in seen
            or transition.get("generation") != generation + 1
            or isinstance(transition.get("generation"), bool)
            or not isinstance(transition.get("observed_at"), str)
            or not transition.get("observed_at")
            or not isinstance(transition.get("reason"), str)
            or not transition.get("reason")
        ):
            raise ValueError("persistent session root thread transition is inconsistent")


def mirror_active_run(directory: Path, session: Mapping[str, Any]) -> None:
    run_id = session.get("current_run_id")
    if not isinstance(run_id, str):
        return
    path = _run_path(directory, run_id)
    run = _read_run_state(path)
    for key in _RUN_MIRROR_FIELDS:
        if key in session:
            run[key] = copy.deepcopy(session[key])
        else:
            run.pop(key, None)
    atomic_write_json(path, run)


def _validate_session_record(path: Path) -> dict[str, Any]:
    data = read_json_object(path, label="session state")
    retired = sorted(RETIRED_SESSION_FIELDS & set(data))
    if retired:
        raise ValueError("persistent session contains retired fields: " + ", ".join(retired))
    if data.get("schema_version") != MMO_SCHEMA_VERSION or isinstance(
        data.get("schema_version"), bool
    ):
        raise ValueError(f"unsupported session state schema; expected {MMO_SCHEMA_VERSION}")
    if data.get("package_version") != PACKAGE_VERSION:
        raise ValueError(f"persistent session package must be {PACKAGE_VERSION}")
    if data.get("profile_version") != PACKAGE_VERSION:
        raise ValueError(f"persistent session profile must be {PACKAGE_VERSION}")
    if data.get("session_kind") not in {"interactive", "noninteractive"}:
        raise ValueError("persistent session kind must be interactive or noninteractive")
    if data.get("root_execution_host") not in ROOT_EXECUTION_HOSTS:
        raise ValueError("persistent session root execution host is invalid")
    _validate_switchyard_identity(data, "persistent session")
    _validate_root_thread_lineage(data)
    socket_path = data.get("root_app_server_socket")
    if not isinstance(socket_path, str) or not Path(socket_path).is_absolute():
        raise ValueError("persistent session app-server socket path is invalid")
    run_sequence = data.get("run_sequence")
    if not isinstance(run_sequence, int) or isinstance(run_sequence, bool) or run_sequence != 1:
        raise ValueError("immutable session run sequence must be exactly 1")
    for key in ("current_run_id", "last_run_id"):
        run_id = data.get(key)
        if run_id is not None and (
            not isinstance(run_id, str) or not SAFE_JOB_ID.fullmatch(run_id)
        ):
            raise ValueError(f"persistent session {key} is invalid")
    expected_id = path.parent.name
    if data.get("session_id") != expected_id:
        raise ValueError(
            "session state identity mismatch: "
            f"directory is {expected_id!r}, record is {data.get('session_id')!r}"
        )
    return data


def _validate_job_record(path: Path) -> dict[str, Any]:
    data = read_json_object(path, label="job state")
    retired = sorted(RETIRED_JOB_FIELDS & set(data))
    if retired:
        raise ValueError("persistent job contains retired fields: " + ", ".join(retired))
    if data.get("schema_version") != MMO_SCHEMA_VERSION or isinstance(
        data.get("schema_version"), bool
    ):
        raise ValueError(f"unsupported job state schema; expected {MMO_SCHEMA_VERSION}")
    if data.get("package_version") != PACKAGE_VERSION:
        raise ValueError(f"persistent job package must be {PACKAGE_VERSION}")
    expected_id = path.parent.name
    if data.get("job_id") != expected_id:
        raise ValueError(
            "job state identity mismatch: "
            f"directory is {expected_id!r}, record is {data.get('job_id')!r}"
        )
    return data


def read_session_record(directory: Path) -> dict[str, Any]:
    """Read and validate the session record owned by ``directory``."""

    return _validate_session_record(session_state_path(directory))


def read_job_record(directory: Path) -> dict[str, Any]:
    """Read and validate the worker-job record owned by ``directory``."""

    return _validate_job_record(job_state_path(directory))


def publish_session_record(
    directory: Path,
    session: Mapping[str, Any],
    *,
    mirror_run: bool,
) -> None:
    """Atomically publish session state and optionally refresh its active-run mirror."""

    atomic_write_json(session_state_path(directory), dict(session))
    if mirror_run:
        mirror_active_run(directory, session)


def publish_initial_session_records(
    directory: Path,
    session: Mapping[str, Any],
    run: Mapping[str, Any],
) -> None:
    """Publish a new run before making its canonical session record visible."""

    run_id = run.get("run_id")
    if not isinstance(run_id, str) or not SAFE_JOB_ID.fullmatch(run_id):
        raise ValueError("initial session run id is invalid")
    if run.get("session_id") != session.get("session_id"):
        raise ValueError("initial session and run identities disagree")
    run_directory = ensure_runs_root(directory) / run_id
    if run_directory.is_symlink():
        raise RuntimeError("initial session run directory cannot be a symlink")
    run_directory.mkdir(mode=0o700)
    atomic_write_json(run_directory / "run.json", dict(run))
    publish_session_record(directory, session, mirror_run=False)


def publish_job_record(directory: Path, job: Mapping[str, Any]) -> None:
    """Atomically publish the canonical worker-job record."""

    atomic_write_json(job_state_path(directory), dict(job))


def terminate_recorded_process_group(
    data: Mapping[str, Any],
    *,
    prefix: str,
    grace_seconds: float,
) -> None:
    pid = data.get(f"{prefix}_pid")
    start_token = data.get(f"{prefix}_start_token")
    if (
        not isinstance(pid, int)
        or isinstance(pid, bool)
        or pid <= 1
        or pid == os.getpid()
        or not isinstance(start_token, str)
    ):
        return
    if bool(data.get(f"{prefix}_process_group_isolated", True)):
        pgid = data.get(f"{prefix}_pgid", pid)
        if not isinstance(pgid, int) or isinstance(pgid, bool) or pgid <= 1 or pgid != pid:
            return
        if process_matches(pid, start_token) or (
            not process_alive(pid) and process_group_alive(pgid)
        ):
            terminate_process_group(pgid, grace_seconds=grace_seconds)
            if process_group_alive(pgid):
                raise RuntimeError(f"{prefix} process group {pgid} did not terminate")
        return
    if process_matches(pid, start_token):
        terminate_process(pid, grace_seconds=grace_seconds)
        if process_matches(pid, start_token):
            raise RuntimeError(f"{prefix} process {pid} did not terminate")


def iter_session_records(*, strict: bool = True) -> list[dict[str, Any]]:
    """Enumerate validated session records without lifecycle side effects."""

    results: list[dict[str, Any]] = []
    for directory in sorted(sessions_root().iterdir(), reverse=True):
        path = session_state_path(directory)
        if not directory.is_dir() or not path.is_file():
            continue
        try:
            results.append(_validate_session_record(path))
        except (OSError, ValueError, json.JSONDecodeError) as exc:
            if strict:
                raise RuntimeError(
                    f"invalid session state blocks safe accounting: {path}: {exc}"
                ) from exc
            continue
    return results


def iter_job_records(*, strict: bool = True) -> list[dict[str, Any]]:
    """Enumerate validated worker-job records without lifecycle side effects."""

    results: list[dict[str, Any]] = []
    for directory in sorted(jobs_root().iterdir(), reverse=True):
        path = job_state_path(directory)
        if not directory.is_dir() or not path.is_file():
            continue
        try:
            results.append(_validate_job_record(path))
        except (OSError, ValueError, json.JSONDecodeError) as exc:
            if strict:
                raise RuntimeError(
                    f"invalid job state blocks safe accounting: {path}: {exc}"
                ) from exc
            continue
    return results


def append_audit(session_id: str, event: str, **data: Any) -> None:
    directory = session_dir(session_id)
    if not isinstance(data.get("run_id"), str):
        with contextlib.suppress(OSError, ValueError, json.JSONDecodeError):
            session = read_session_record(directory)
            run_id = session.get("current_run_id") or session.get("last_run_id")
            if isinstance(run_id, str):
                data["run_id"] = run_id
    append_jsonl(
        directory / "audit.jsonl",
        {"timestamp": utc_now(), "event": event, "session_id": session_id, **data},
    )
