"""Create a Modal VM Sandbox whose /workspace is an AIStor Memory Bucket.

Required host configuration:
  * Modal authentication (`modal setup` or MODAL_TOKEN_ID/MODAL_TOKEN_SECRET)
  * AIMEM_BUCKET and AIMEM_ENDPOINT
  * a Modal Secret (default: aimem-creds) containing
    AIMEM_ACCESS_KEY and AIMEM_SECRET_KEY

Run:
    uv run python usage.py
"""

from __future__ import annotations

import base64
import binascii
import os
import re
import sys
import threading
import time
import uuid
from pathlib import Path
from urllib.parse import urlparse

import modal


HERE = Path(__file__).resolve().parent
CONTRACT_DIR = HERE.parent / "agent-worker-runtime"
sys.path.insert(0, str(CONTRACT_DIR))
from contract import (  # noqa: E402
    RemoteCommandWorkspaceTools,
    emit_trace,
    run_contract,
    without_storage_credentials,
)

APP_NAME = os.environ.get("AIMEM_MODAL_APP", "aimem-workspace")
SECRET_NAME = os.environ.get("AIMEM_MODAL_SECRET", "aimem-creds")
MOUNTPOINT = os.environ.get("AIMEM_MOUNTPOINT", "/workspace")
LOCAL = os.environ.get("AIMEM_LOCAL", "/var/lib/aimem-local")
VERSION = os.environ.get("AIMEM_VERSION", "latest")
LOCAL_BINARY = os.environ.get("AIMEM_BIN")
DL_BASE = os.environ.get(
    "AIMEM_DL_BASE", "https://dl.min.io/aistor/aimem/release"
).rstrip("/")
MOUNT_NOT_VISIBLE = 98


def required_env(name: str) -> str:
    value = os.environ.get(name)
    if not value:
        raise RuntimeError(f"{name} is required; copy .env.example and fill it in")
    return value


def validate_absolute_path(name: str, value: str) -> str:
    path = Path(value)
    if not path.is_absolute() or ".." in path.parts:
        raise RuntimeError(f"{name} must be an absolute, normalized path")
    return value


def validate_endpoint(value: str) -> str:
    parsed = urlparse(value)
    if parsed.scheme != "https" or not parsed.hostname:
        raise RuntimeError("AIMEM_ENDPOINT must be an HTTPS URL with a host")
    return value


def build_image() -> modal.Image:
    image = (
        modal.Image.from_registry("ubuntu:24.04")
        .apt_install("ca-certificates", "coreutils", "curl", "util-linux")
        .add_local_file(
            HERE / "aimem-mount",
            "/usr/local/bin/aimem-mount",
            copy=True,
        )
        .run_commands(
            "chmod 0755 /usr/local/bin/aimem-mount",
            "mkdir -p /workspace /var/lib/aimem-local",
        )
    )

    if LOCAL_BINARY:
        binary = Path(LOCAL_BINARY).expanduser().resolve()
        if not binary.is_file():
            raise RuntimeError(f"AIMEM_BIN is not a file: {binary}")
        return image.add_local_file(
            binary,
            "/usr/local/bin/aimem",
            copy=True,
        ).run_commands("chmod 0755 /usr/local/bin/aimem")

    # Values enter a shell command during image construction. Accept release
    # tags and "latest", not arbitrary shell text.
    if not re.fullmatch(r"[A-Za-z0-9._:+-]+", VERSION):
        raise RuntimeError(f"invalid AIMEM_VERSION: {VERSION!r}")
    if not re.fullmatch(r"https://[A-Za-z0-9./_:+-]+", DL_BASE):
        raise RuntimeError("AIMEM_DL_BASE must be a simple HTTPS URL")

    # latest is served under the bare file name; a pinned tag is a suffix.
    suffix = "" if VERSION == "latest" else f".{VERSION}"
    url = f"{DL_BASE}/linux-amd64/aimem{suffix}"
    return image.run_commands(
        f"curl -fsSL {url} -o /usr/local/bin/aimem",
        f"curl -fsSL {url}.sha256sum -o /tmp/aimem.sha256sum",
        "echo \"$(cut -d' ' -f1 /tmp/aimem.sha256sum)  /usr/local/bin/aimem\" | sha256sum -c -",
        "chmod 0755 /usr/local/bin/aimem",
        "rm -f /tmp/aimem.sha256sum",
        "/usr/local/bin/aimem --version",
    )


def sandbox_env() -> dict[str, str]:
    mountpoint = validate_absolute_path("AIMEM_MOUNTPOINT", MOUNTPOINT)
    local = validate_absolute_path("AIMEM_LOCAL", LOCAL)
    if (
        local == mountpoint
        or Path(mountpoint) in Path(local).parents
        or Path(local) in Path(mountpoint).parents
    ):
        raise RuntimeError("AIMEM_LOCAL and AIMEM_MOUNTPOINT must not overlap")

    values = {
        "AIMEM_BUCKET": required_env("AIMEM_BUCKET"),
        "AIMEM_ENDPOINT": validate_endpoint(required_env("AIMEM_ENDPOINT")),
        "AIMEM_MOUNTPOINT": mountpoint,
        "AIMEM_LOCAL": local,
        "AIMEM_AGENT": os.environ.get("AIMEM_AGENT", "auto"),
    }
    # Signing credentials, including AIMEM_SESSION_TOKEN, belong exclusively in
    # the Modal Secret passed to Sandbox.create().
    for name in ("AIMEM_REGION", "AIMEM_PREFIX"):
        if value := os.environ.get(name):
            values[name] = value
    return values


def create_sandbox(app: modal.App, image: modal.Image) -> modal.Sandbox:
    return modal.Sandbox.create(
        "/usr/local/bin/aimem-mount",
        app=app,
        image=image,
        env=sandbox_env(),
        secrets=[modal.Secret.from_name(SECRET_NAME)],
        timeout=10 * 60,
        cpu=2,
        memory=4096,
        experimental_options={"vm_runtime": True},
    )


def discard_sandbox_output(sb: modal.Sandbox) -> list[threading.Thread]:
    """Drain privileged mount logs without copying them into the host output."""

    def discard(stream: object) -> None:
        try:
            stream.read()  # type: ignore[attr-defined]
        except Exception:
            # Sandbox termination can close a stream while the reader is
            # blocked. Cleanup must not obscure the smoke-test result.
            pass

    readers = [
        threading.Thread(target=discard, args=(stream,), daemon=True)
        for stream in (sb.stdout, sb.stderr)
    ]
    for reader in readers:
        reader.start()
    return readers


def wait_for_mount(sb: modal.Sandbox) -> None:
    for _ in range(60):
        fuse = sb.exec("test", "-c", "/dev/fuse", timeout=5)
        fuse.wait()
        process = sb.exec("mountpoint", "-q", MOUNTPOINT, timeout=5)
        process.wait()
        if fuse.returncode == 0 and process.returncode == 0:
            return
        if sb.poll() is not None:
            # Do not echo the mount process's output: it runs with the signing
            # Secret in its environment and upstream diagnostics are not part
            # of our credential-redaction contract.
            raise RuntimeError("AIMem exited before mounting; inspect protected Modal logs")
        time.sleep(1)
    raise RuntimeError(f"AIMem did not mount {MOUNTPOINT} within 60 seconds")


def run_process(
    sb: modal.Sandbox, *command: str, timeout: int | None = None
) -> tuple[int, str, str]:
    process = sb.exec(*command, timeout=timeout)
    output = {"stdout": "", "stderr": ""}

    def drain(name: str) -> None:
        try:
            output[name] = getattr(process, name).read()
        except Exception:
            # Modal may close streams during timeout or Sandbox termination.
            # Preserve the process result rather than masking it with KeyError.
            output[name] = ""

    # Drain both pipes while the command runs so a verbose child cannot block
    # forever on a full stdout or stderr buffer.
    readers = [
        threading.Thread(target=drain, args=(name,), daemon=True)
        for name in ("stdout", "stderr")
    ]
    for reader in readers:
        reader.start()
    process.wait()
    for reader in readers:
        reader.join()
    return process.returncode, output["stdout"], output["stderr"]


def run_checked(sb: modal.Sandbox, *command: str) -> str:
    returncode, stdout, stderr = run_process(sb, *command)
    if returncode != 0:
        raise RuntimeError(
            f"sandbox command failed ({returncode}); "
            f"stderr redacted ({len(stderr.encode())} bytes)"
        )
    return stdout


BROKER_SCRIPT = """\
for attempt in $(seq 1 40); do
    mountpoint -q -- "$1" && break
    test "$attempt" -eq 40 && exit 98
    sleep 0.25
done
while IFS= read -r request; do
    command=$(printf '%s' "$request" | base64 -d) || exit 99
    stdout_file=$(mktemp) || exit 99
    stderr_file=$(mktemp) || { rm -f "$stdout_file"; exit 99; }
    bash -lc "$command" >"$stdout_file" 2>"$stderr_file"
    status=$?
    printf '%s\t' "$status"
    base64 -w0 "$stdout_file"
    printf '\t'
    base64 -w0 "$stderr_file"
    printf '\n'
    rm -f "$stdout_file" "$stderr_file"
done
"""


class ModalCommandBroker:
    """Run every workspace command in one mount-owning Modal exec namespace."""

    def __init__(self, sb: modal.Sandbox) -> None:
        self.process = sb.exec(
            "bash",
            "-lc",
            BROKER_SCRIPT,
            "aimem-command-broker",
            MOUNTPOINT,
            timeout=10 * 60,
            bufsize=1,
        )
        self.stdout = iter(self.process.stdout)
        self.lock = threading.Lock()
        self.stderr_reader = threading.Thread(
            target=self._discard_stderr,
            daemon=True,
        )
        self.stderr_reader.start()

    def _discard_stderr(self) -> None:
        try:
            self.process.stderr.read()
        except Exception:
            pass

    def __call__(self, command: str, timeout: int) -> tuple[int, str, str]:
        del timeout  # The shared command already carries its own GNU timeout.
        request = base64.b64encode(command.encode()).decode() + "\n"
        with self.lock:
            self.process.stdin.write(request)
            self.process.stdin.drain()
            try:
                response = next(self.stdout).rstrip("\n")
            except StopIteration as exc:
                self.process.wait()
                raise RuntimeError(
                    f"Modal command broker exited ({self.process.returncode})"
                ) from exc
        fields = response.split("\t")
        if len(fields) != 3:
            raise RuntimeError("Modal command broker returned an invalid response")
        status, stdout, stderr = fields
        try:
            return (
                int(status),
                base64.b64decode(stdout, validate=True).decode(),
                base64.b64decode(stderr, validate=True).decode(),
            )
        except (ValueError, binascii.Error, UnicodeDecodeError) as exc:
            raise RuntimeError(
                "Modal command broker returned an invalid encoded response"
            ) from exc

    def close(self) -> None:
        try:
            self.process.stdin.write_eof()
            self.process.stdin.drain()
            self.process.wait()
        finally:
            self.stderr_reader.join(timeout=5)


class ModalWorkspaceTools(RemoteCommandWorkspaceTools):
    def __init__(self, sb: modal.Sandbox) -> None:
        self.broker = ModalCommandBroker(sb)
        super().__init__(
            sb.object_id,
            self.broker,
            workspace=MOUNTPOINT,
        )

    def close(self) -> None:
        self.broker.close()


def workspace_tools(sb: modal.Sandbox) -> ModalWorkspaceTools:
    return ModalWorkspaceTools(
        sb,
    )


def verify_installed_skill(sb: modal.Sandbox) -> None:
    run_checked(
        sb,
        "bash",
        "-lc",
        'for d in "$HOME/.agents/skills/aimem" "$HOME/.claude/skills/aimem"; do '
        'test -s "$d/SKILL.md" || exit 1; '
        "grep -q '^name: aimem$' \"$d/SKILL.md\" || exit 1; "
        "for r in navigation workspace memory metadata annotations; do "
        'test -s "$d/references/$r.md" || exit 1; '
        "done; "
        "done",
    )


def main() -> None:
    app = modal.App.lookup(APP_NAME, create_if_missing=True)
    image = build_image()
    marker = f"external-worker/from-modal-{uuid.uuid4().hex}.txt"
    value = f"hello from modal {uuid.uuid4()}\n"

    with modal.enable_output():
        first = create_sandbox(app, image)
    first_readers = discard_sandbox_output(first)
    first_tools: ModalWorkspaceTools | None = None
    try:
        wait_for_mount(first)
        first_tools = workspace_tools(first)
        verify_installed_skill(first)
        # The trusted worker keeps model/storage credentials out of its own
        # environment and reaches the sandbox only through typed workspace
        # tools. The Sandbox alone owns the AIMem signing Secret.
        with without_storage_credentials():
            run_contract(first_tools, marker, value)
        print(f"wrote {MOUNTPOINT}/{marker}: {value.strip()}")
    finally:
        try:
            if first_tools is not None:
                emit_trace(first_tools)
                first_tools.close()
        finally:
            first.terminate(wait=True)
            for reader in first_readers:
                reader.join(timeout=5)

    with modal.enable_output():
        second = create_sandbox(app, image)
    second_readers = discard_sandbox_output(second)
    try:
        wait_for_mount(second)
        second_tools = workspace_tools(second)
        with without_storage_credentials():
            run_contract(second_tools, marker, value, verify_only=True)
        print(f"remount verified {MOUNTPOINT}/{marker}: {value.strip()}")
        second_tools.remove_tree(marker)
    finally:
        try:
            if "second_tools" in locals():
                emit_trace(second_tools)
                second_tools.close()
        finally:
            second.terminate(wait=True)
            for reader in second_readers:
                reader.join(timeout=5)

    print("MODAL VM SANDBOX AIMEM SMOKE OK")


if __name__ == "__main__":
    main()