"""Read-only Phase 2 runtime introspection with a bounded SSH transport.

The wrapper is deliberately sent as ``subprocess.run(input=payload)``.  No
JSON, ABI data, or base64 representation is placed in the Windows SSH argv;
the default argv is the fixed, short command documented by
``DEFAULT_REMOTE_COMMAND``.

The remote wrapper reads the current issuer unit's MainPID, WorkingDirectory,
and PYTHONPATH, then resolves only the named importable modules.  It does not
walk release trees, ``/opt``, historical worktrees, SOUL, or MEMORY, and it
does not call a database, signer, Bridge, Task API, or mutation endpoint.

This module is locally testable, but the module import and systemd/POSIX
behaviour still require external Linux live validation.  Phase 2 status is
always explicit about all mutation classes being false.
"""

from __future__ import annotations

import argparse
import ast
import hashlib
import json
import subprocess
import sys
from pathlib import Path
from typing import Any, Callable, Mapping, Sequence


DEFAULT_SSH = "ssh.exe"
DEFAULT_HOST = "hw"
DEFAULT_REMOTE_COMMAND = 'test "$(hostname)" = cn-network-mkt && sudo -n python3 -'
DEFAULT_TIMEOUT_SECONDS = 60
DEFAULT_ISSUER_UNIT = "aios-task-principal-issuer.service"
DEFAULT_ADMISSION_MODULE = "aios_task_principal_admission"
DEFAULT_MIGRATION_MODULE = "aios_task_principal_storage_migration"
DEFAULT_REQUIRED_ADMISSION_FUNCTIONS = (
    "offline_migrate_task_principal_schema",
    "supersede_p0_scope",
)
DEFAULT_MIGRATION_VERIFIER = "SSHOwnerBootstrapVerifier"
MAX_ARGV_BYTES = 512
MAX_REQUEST_BYTES = 32 * 1024

MUTATION_FLAGS = (
    "task_db_mutation",
    "scope_mutation",
    "git_mutation",
    "service_mutation",
)


class RemoteTransportError(RuntimeError):
    """A safe, non-secret description of an SSH transport failure."""


class Phase2StatusError(RuntimeError):
    """The remote result was not a valid fail-closed Phase2 status."""


def build_ssh_argv(
    ssh: str = DEFAULT_SSH,
    host: str = DEFAULT_HOST,
) -> list[str]:
    """Return the only SSH argv shape permitted by the Phase2 transport."""

    argv = [
        ssh,
        "-T",
        "-o",
        "BatchMode=yes",
        "-o",
        "ConnectTimeout=12",
        host,
        DEFAULT_REMOTE_COMMAND,
    ]
    argv_bytes = sum(len(part.encode("utf-8")) + 1 for part in argv)
    if argv_bytes > MAX_ARGV_BYTES:
        raise ValueError("SSH argv exceeds the Phase2 hard limit")
    return argv


def run_remote_python(
    payload: bytes,
    *,
    ssh: str = DEFAULT_SSH,
    host: str = DEFAULT_HOST,
    timeout: float = DEFAULT_TIMEOUT_SECONDS,
    runner: Callable[..., Any] = subprocess.run,
) -> bytes:
    """Send *payload* as stdin to the fixed remote Python command.

    Error messages intentionally contain no stderr text.  SSH stderr can
    include command output, environment values, or credentials emitted by a
    misconfigured remote service.
    """

    if not isinstance(payload, bytes):
        raise TypeError("payload must be bytes")
    if timeout <= 0:
        raise ValueError("timeout must be positive")

    argv = build_ssh_argv(ssh=ssh, host=host)
    try:
        completed = runner(
            argv,
            input=payload,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            timeout=timeout,
            check=False,
        )
    except subprocess.TimeoutExpired as exc:
        raise RemoteTransportError(
            f"remote SSH timed out after {timeout:g} seconds"
        ) from exc
    except OSError as exc:
        raise RemoteTransportError("remote SSH could not be started") from exc

    return_code = getattr(completed, "returncode", None)
    if return_code != 0:
        raise RemoteTransportError(f"remote SSH exited with status {return_code}")
    output = getattr(completed, "stdout", b"")
    if isinstance(output, str):
        output = output.encode("utf-8")
    if not isinstance(output, bytes):
        raise RemoteTransportError("remote SSH returned a non-byte stdout value")
    return output


def module_sha256(path: str | Path) -> str:
    """Hash one resolved module file without searching neighbouring files."""

    digest = hashlib.sha256()
    with Path(path).open("rb") as source:
        for block in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def validate_public_function_ast(
    source: str,
    required_functions: Sequence[str] = DEFAULT_REQUIRED_ADMISSION_FUNCTIONS,
) -> dict[str, Any]:
    """Validate public symbols and the source-level export mechanism.

    The live admission overlay intentionally re-exports the legacy public API
    through a small ``globals()[name] = value`` loop.  Requiring every public
    function to be a direct ``def`` would therefore reject the exact deployed
    module even though the runtime symbol and signature are present.  The AST
    check still proves that the module has either a direct/imported binding or
    an explicit dynamic re-export mechanism; the caller separately verifies
    the runtime callable and its signature.
    """

    try:
        tree = ast.parse(source)
    except SyntaxError as exc:
        return {
            "ok": False,
            "missing": list(required_functions),
            "error": f"AST_PARSE_ERROR:{exc.msg}",
        }

    top_level = {
        node.name: node
        for node in tree.body
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
        and not node.name.startswith("_")
    }
    imported = {
        alias.asname or alias.name.split(".", 1)[0]
        for node in tree.body
        if isinstance(node, (ast.Import, ast.ImportFrom))
        for alias in node.names
        if alias.name != "*"
    }
    dynamic_reexport = any(
        isinstance(node, ast.Call)
        and isinstance(node.func, ast.Name)
        and node.func.id == "globals"
        for node in ast.walk(tree)
    )
    bound = set(top_level) | imported
    missing = [
        name
        for name in required_functions
        if name not in bound and not dynamic_reexport
    ]
    return {
        "ok": not missing,
        "missing": missing,
        "functions": sorted(top_level),
        "imported": sorted(imported),
        "dynamic_reexport": dynamic_reexport,
    }


def make_phase2_status(
    status: str,
    errors: Sequence[str] = (),
    **details: Any,
) -> dict[str, Any]:
    """Create a status object whose mutation flags are always false."""

    result: dict[str, Any] = {
        "kind": "phase2-status",
        "version": 1,
        "status": status,
        "read_only": True,
        "fail_closed": status != "PASS",
    }
    result.update({flag: False for flag in MUTATION_FLAGS})
    if errors:
        result["errors"] = list(errors)
    result.update(details)
    return result


def validate_phase2_status(status: Mapping[str, Any]) -> dict[str, Any]:
    """Validate the result contract and reject any mutation claim."""

    if status.get("kind") != "phase2-status":
        raise Phase2StatusError("remote result is not phase2-status")
    for flag in MUTATION_FLAGS:
        if status.get(flag) is not False:
            raise Phase2StatusError(f"remote result has invalid {flag}")
    if status.get("status") not in {"PASS", "FAIL"}:
        raise Phase2StatusError("remote result has invalid status")
    return dict(status)


def _json_request(request: Mapping[str, Any]) -> str:
    """Serialize only a small request into the stdin-delivered wrapper."""

    encoded = json.dumps(dict(request), ensure_ascii=True, separators=(",", ":"))
    if len(encoded.encode("utf-8")) > MAX_REQUEST_BYTES:
        raise ValueError("Phase2 request is unexpectedly large")
    return encoded


REMOTE_SCRIPT_PREFIX = r'''# Phase2 read-only runtime wrapper.  Do not add mutations here.
import ast
import hashlib
import importlib
import importlib.util
import inspect
import json
import os
import shlex
import subprocess
import sys
from pathlib import Path

sys.dont_write_bytecode = True
REQUEST = json.loads('''
REMOTE_SCRIPT_SUFFIX = r''')

MUTATION_FLAGS = ("task_db_mutation", "scope_mutation", "git_mutation", "service_mutation")

def status(status_name, errors=(), **details):
    result = {
        "kind": "phase2-status",
        "version": 1,
        "status": status_name,
        "read_only": True,
        "fail_closed": status_name != "PASS",
    }
    result.update({flag: False for flag in MUTATION_FLAGS})
    if errors:
        result["errors"] = list(errors)
    result.update(details)
    return result

def systemd_value(unit, property_name):
    try:
        completed = subprocess.run(
            ["systemctl", "show", unit, "--property=" + property_name, "--value", "--no-pager"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            check=False,
            timeout=12,
        )
    except (OSError, subprocess.TimeoutExpired):
        raise RuntimeError("SYSTEMD_QUERY_FAILED")
    if completed.returncode != 0:
        raise RuntimeError("SYSTEMD_QUERY_FAILED")
    return completed.stdout.decode("utf-8", "strict").strip()

def parse_pythonpath(environment_value):
    for token in shlex.split(environment_value):
        if token.startswith("PYTHONPATH="):
            return token[len("PYTHONPATH="):].split(os.pathsep)
    return []

def runtime_identity():
    unit = REQUEST.get("issuer_unit", "aios-task-principal-issuer.service")
    pid = REQUEST.get("issuer_pid")
    working_directory = REQUEST.get("working_directory")
    pythonpath = REQUEST.get("pythonpath")
    # Regression callers may pin one identity component while obtaining the
    # remaining components from the live unit.  In particular, never replace
    # an explicitly supplied PID with a newly queried MainPID.
    if pid is None:
        pid = int(systemd_value(unit, "MainPID"))
    if working_directory is None:
        working_directory = systemd_value(unit, "WorkingDirectory")
    if pythonpath is None:
        pythonpath = parse_pythonpath(systemd_value(unit, "Environment"))
    if not pid or not working_directory or working_directory == "-":
        raise RuntimeError("ISSUER_RUNTIME_IDENTITY_INCOMPLETE")
    if isinstance(pythonpath, str):
        pythonpath = pythonpath.split(os.pathsep)
    return {
        "unit": unit,
        "main_pid": int(pid),
        "working_directory": working_directory,
        "pythonpath": [item for item in pythonpath if item],
    }

def setup_import_path(identity):
    path = [identity["working_directory"]]
    path.extend(identity["pythonpath"])
    path.extend(item for item in sys.path if item and item not in path)
    sys.path[:] = path

def importable_module(module_name):
    try:
        spec = importlib.util.find_spec(module_name)
        if spec is None or not spec.origin or spec.origin in {"built-in", "frozen"}:
            raise RuntimeError("MODULE_NOT_IMPORTABLE:" + module_name)
        module = importlib.import_module(module_name)
    except Exception as exc:
        if isinstance(exc, RuntimeError):
            raise
        raise RuntimeError("MODULE_IMPORT_FAILED:" + module_name) from exc
    module_path = getattr(module, "__file__", None)
    if not module_path:
        raise RuntimeError("MODULE_HAS_NO_FILE:" + module_name)
    return module, Path(module_path).resolve()

def hash_file(path):
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for block in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()

def inspect_admission(module, path, required, expected_hash):
    source = path.read_text(encoding="utf-8")
    tree = ast.parse(source)
    nodes = {
        node.name: node for node in tree.body
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
        and not node.name.startswith("_")
    }
    imported = {
        alias.asname or alias.name.split(".", 1)[0]
        for node in tree.body
        if isinstance(node, (ast.Import, ast.ImportFrom))
        for alias in node.names
        if alias.name != "*"
    }
    dynamic_reexport = any(
        isinstance(node, ast.Call)
        and isinstance(node.func, ast.Name)
        and node.func.id == "globals"
        for node in ast.walk(tree)
    )
    bound = set(nodes) | imported
    missing = [
        name
        for name in required
        if name not in bound and not dynamic_reexport
    ]
    if missing:
        raise RuntimeError("PUBLIC_SYMBOL_MISSING:" + ",".join(missing))
    signatures = {}
    for name in required:
        value = getattr(module, name, None)
        if value is None:
            raise RuntimeError("PUBLIC_SYMBOL_NOT_IMPORTABLE:" + name)
        try:
            signatures[name] = str(inspect.signature(value))
        except (TypeError, ValueError):
            raise RuntimeError("SIGNATURE_UNAVAILABLE:" + name)
    actual_hash = hash_file(path)
    if expected_hash and actual_hash != expected_hash:
        raise RuntimeError("MODULE_HASH_MISMATCH:admission")
    return {
        "name": module.__name__,
        "path": str(path),
        "sha256": actual_hash,
        "signatures": signatures,
        "ast_functions": sorted(nodes),
        "ast_imported": sorted(imported),
        "ast_dynamic_reexport": dynamic_reexport,
    }

def inspect_migration(module, path, verifier_name, expected_hash):
    verifier = getattr(module, verifier_name, None)
    if verifier is None:
        raise RuntimeError("MIGRATION_VERIFIER_MISSING:" + verifier_name)
    try:
        signature = str(inspect.signature(verifier))
    except (TypeError, ValueError):
        raise RuntimeError("MIGRATION_VERIFIER_SIGNATURE_UNAVAILABLE")
    actual_hash = hash_file(path)
    if expected_hash and actual_hash != expected_hash:
        raise RuntimeError("MODULE_HASH_MISMATCH:migration")
    return {
        "name": module.__name__,
        "path": str(path),
        "sha256": actual_hash,
        "verifier": verifier_name,
        "signature": signature,
    }

def main():
    try:
        identity = runtime_identity()
        setup_import_path(identity)
        admission_name = REQUEST.get("admission_module", "aios_task_principal_admission")
        migration_name = REQUEST.get("migration_module", "aios_task_principal_storage_migration")
        if migration_name == admission_name:
            raise RuntimeError("MIGRATION_MODULE_NOT_DISTINCT")
        required = tuple(REQUEST.get("required_admission_functions", (
            "offline_migrate_task_principal_schema", "supersede_p0_scope")))
        admission, admission_path = importable_module(admission_name)
        migration, migration_path = importable_module(migration_name)
        admission_info = inspect_admission(
            admission,
            admission_path,
            required,
            REQUEST.get("expected_admission_sha256"),
        )
        migration_info = inspect_migration(
            migration,
            migration_path,
            REQUEST.get("migration_verifier", "SSHOwnerBootstrapVerifier"),
            REQUEST.get("expected_migration_sha256"),
        )
        print(json.dumps(status(
            "PASS",
            issuer=identity,
            modules={"admission": admission_info, "migration": migration_info},
            ast_public_functions=list(required),
        ), sort_keys=True, separators=(",", ":")))
    except Exception as exc:
        # Emit only stable error classes and never exception text from imports.
        error = str(exc).split(":", 1)[0] or "PHASE2_INTROSPECTION_FAILED"
        print(json.dumps(status("FAIL", errors=[error]), sort_keys=True, separators=(",", ":")))

if __name__ == "__main__":
    main()
'''


def build_remote_payload(request: Mapping[str, Any] | None = None) -> bytes:
    """Build the stdin wrapper; request data never enters SSH argv."""

    request = {} if request is None else request
    # repr() makes the embedded JSON a valid Python string for every request
    # value, including apostrophes and backslashes; it still remains stdin.
    request_literal = repr(_json_request(request))
    return (REMOTE_SCRIPT_PREFIX + request_literal + REMOTE_SCRIPT_SUFFIX).encode("utf-8")


def _load_request(args: argparse.Namespace) -> dict[str, Any]:
    request: dict[str, Any] = {}
    if args.request_json:
        loaded = json.loads(args.request_json)
        if not isinstance(loaded, dict):
            raise ValueError("request JSON must be an object")
        request.update(loaded)
    if args.request_file:
        loaded = json.loads(Path(args.request_file).read_text(encoding="utf-8"))
        if not isinstance(loaded, dict):
            raise ValueError("request JSON must be an object")
        request.update(loaded)

    overrides = {
        "issuer_unit": args.issuer_unit,
        "issuer_pid": args.issuer_pid,
        "working_directory": args.working_directory,
        "pythonpath": args.pythonpath,
        "admission_module": args.admission_module,
        "migration_module": args.migration_module,
        "expected_admission_sha256": args.expected_admission_sha256,
        "expected_migration_sha256": args.expected_migration_sha256,
    }
    request.update({key: value for key, value in overrides.items() if value is not None})
    return request


def _write_status(path: str | Path, status: Mapping[str, Any]) -> None:
    destination = Path(path)
    destination.write_text(
        json.dumps(dict(status), ensure_ascii=True, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )


def main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--request-json")
    parser.add_argument("--request-file")
    parser.add_argument("--output", default="phase2-status.json")
    parser.add_argument("--ssh", default=DEFAULT_SSH)
    parser.add_argument("--host", default=DEFAULT_HOST)
    parser.add_argument("--timeout", type=float, default=DEFAULT_TIMEOUT_SECONDS)
    parser.add_argument("--issuer-unit")
    parser.add_argument("--issuer-pid", type=int)
    parser.add_argument("--working-directory")
    parser.add_argument("--pythonpath", nargs="*")
    parser.add_argument("--admission-module")
    parser.add_argument("--migration-module")
    parser.add_argument("--expected-admission-sha256")
    parser.add_argument("--expected-migration-sha256")
    args = parser.parse_args(argv)

    try:
        request = _load_request(args)
        payload = build_remote_payload(request)
        stdout = run_remote_python(
            payload,
            ssh=args.ssh,
            host=args.host,
            timeout=args.timeout,
        )
        remote_status = json.loads(stdout.decode("utf-8"))
        status = validate_phase2_status(remote_status)
    except (OSError, ValueError, json.JSONDecodeError, RemoteTransportError, Phase2StatusError) as exc:
        status = make_phase2_status("FAIL", errors=[str(exc).split(":", 1)[0]])
        _write_status(args.output, status)
        return 1

    _write_status(args.output, status)
    return 0 if status["status"] == "PASS" else 1


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