Store and read back memory on NAMS

"""Private, restartable state for the hosted REST tutorials; no I/O at import."""

from __future__ import annotations

import hashlib
import json
import os
import re
import tempfile
from datetime import datetime, timezone
from pathlib import Path
from urllib.parse import urlsplit, urlunsplit
from uuid import uuid4


def read_state(path):
    """Read a local ledger without credentials, SDK imports, or service access."""
    path = Path(path)
    if path.is_symlink():
        raise RuntimeError("Refusing a symlink for tutorial state")
    data = json.loads(path.read_text())
    if not isinstance(data, dict) or data.get("schema_version") != 1:
        raise RuntimeError("Unsupported or incomplete tutorial state")
    if not isinstance(data.get("identity"), dict):
        raise RuntimeError("Missing tutorial identity")
    if (
        not isinstance(data.get("resources"), dict)
        or not data.get("run_token")
        or "pending" not in data
    ):
        raise RuntimeError("Incomplete tutorial state")
    if any(not isinstance(entries, dict) for entries in data["resources"].values()):
        raise RuntimeError("Invalid resource ledger")
    if any(
        not isinstance(entry, dict)
        for entries in data["resources"].values()
        for entry in entries.values()
    ):
        raise RuntimeError("Invalid resource disposition")
    datetime.fromisoformat(data["started_at"])
    return data


def new_run_status(data):
    """Classify recorded evidence only; never infer absence from an empty ledger."""
    if data.get("pending"):
        return "needs_reconciliation"
    resources = data["resources"]
    entries = [entry for group in resources.values() for entry in group.values()]
    if (
        data.get("remote_write_started") is False
        and not entries
        and not any(data.get(key) for key in ("conversation_id", "run_id", "skill_id"))
        and not data.get("ontology", {}).get("clone_id")
    ):
        return "no_remote_write_started"
    disposed = data.get("cleanup_complete") is True or (
        data.get("ontology", {}).get("status") == "complete"
        and data["ontology"].get("restored") is True
        and data["ontology"].get("deleted") is True
    )
    if (
        disposed
        and entries
        and all(entry.get("status") in {"deleted", "deleted_with_ontology"} for entry in entries)
    ):
        return "recorded_disposition_complete"
    return "owner_disposition_required"


# Restoration target and exercise outcome shown by inspect-file.
ONTOLOGY_SUMMARY_KEYS = (
    "previous",
    "strict",
    "clone_id",
    "template",
    "status",
    "restored",
    "deleted",
    "exercise_error",
)


def inspect_file(path):
    """Return a redacted local summary; this never authenticates or enables writes."""
    data = read_state(path)
    endpoint = urlsplit(str(data["identity"].get("endpoint", "")))
    # Do not expose credentials/query parameters even from a manually altered file.
    endpoint = urlunsplit(
        (endpoint.scheme, endpoint.netloc.rsplit("@", 1)[-1], endpoint.path, "", "")
    )
    summary = {
        "local_only": True,
        "service_checked": False,
        "lesson": data.get("lesson"),
        "run_token": data["run_token"],
        "started_at": data["started_at"],
        "identity": {
            "endpoint": endpoint,
            "workspace_id": data["identity"].get("workspace_id"),
        },
        "operator_context": data.get("operator_context", {}),
        "pending": data.get("pending"),
        "new_run_status": new_run_status(data),
        "resources": {
            kind: {
                resource_id: {
                    key: entry[key] for key in ("status", "owned", "disposition") if key in entry
                }
                for resource_id, entry in entries.items()
            }
            for kind, entries in data["resources"].items()
        },
    }
    ontology = data.get("ontology")
    if isinstance(ontology, dict):
        # The owner needs the saved prior binding to restore the workspace.
        # Schema documents are omitted; the IDs identify them exactly.
        binding_keys = ("version_id", "ontology_id", "revision", "validation_mode")
        summary["ontology"] = {
            key: (
                {k: value[k] for k in binding_keys if k in value}
                if isinstance(value, dict)
                else value
            )
            for key, value in ontology.items()
            if key in ONTOLOGY_SUMMARY_KEYS
        }
    return summary


def check_new_run(path, next_state):
    """Permit a separate path only after no-write proof or recorded full disposal."""
    status = new_run_status(read_state(path))
    if status not in {"no_remote_write_started", "recorded_disposition_complete"}:
        raise RuntimeError(
            "No automatic new run: uncertain, retained, or legacy state needs owner disposition"
        )
    next_state = Path(next_state)
    if next_state.exists() or next_state.is_symlink():
        raise RuntimeError("Choose an unused state path; the previous run must remain intact")
    if next_state.parent.resolve() == Path(path).parent.resolve():
        raise RuntimeError("Use a separate run directory so reports and archives cannot collide")
    if next_state.parent.exists() and any(next_state.parent.iterdir()):
        raise RuntimeError("Choose a new or empty run directory; preserve its existing reports")
    return status


def digest(text: str) -> str:
    return hashlib.sha256(text.encode()).hexdigest()


def connection(settings):
    """Use the same resolved NamsSettings for SDK and tutorial HTTP requests."""
    config = settings.nams
    parts = urlsplit(str(config.endpoint).rstrip("/"))
    if parts.scheme not in {"http", "https"} or not parts.netloc:
        raise ValueError("A hosted HTTP(S) endpoint is required")
    if parts.username or parts.password or parts.query or parts.fragment:
        raise ValueError("Endpoint must not contain credentials, a query, or a fragment")
    if config.transport_mode == "bridge" or (
        config.transport_mode == "auto" and not re.search(r"/v\d+(?:/|$)", parts.path)
    ):
        raise ValueError("These tutorials require the hosted REST API")
    endpoint = urlunsplit((parts.scheme.lower(), parts.netloc.lower(), parts.path, "", ""))
    key = config.api_key.get_secret_value() if config.api_key else ""
    if not key:
        raise ValueError("Set MEMORY_API_KEY before starting the tutorial")
    # Tutorial identity must not differ between the SDK and direct HTTP calls.
    headers = dict(config.headers)
    if any(k.lower() in {"authorization", "x-workspace-id"} for k in headers):
        raise ValueError("Use api_key/workspace_id settings, not overriding auth/workspace headers")
    headers["Authorization"] = f"Bearer {key}"
    if config.workspace_id:
        headers["X-Workspace-Id"] = config.workspace_id
    resolved = {"endpoint": endpoint, "workspace_id": config.workspace_id}
    return resolved, headers, config.timeout, key


def credential_digest(key: str, salt: str) -> str:
    """Bind state to the credential that seeded it, without keeping a directly
    comparable digest of that credential on disk. A per-run salt and a
    deliberately slow KDF mean the stored value answers one question -- "is this
    the same key as last time?" -- and is not worth attacking offline. Message
    content keeps ``digest``: it is a change detector, not a secret."""
    return hashlib.scrypt(
        key.encode(), salt=bytes.fromhex(salt), n=16384, r=8, p=1, maxmem=64 * 1024 * 1024, dklen=32
    ).hex()


def identity(settings, salt=None):
    resolved, _headers, _timeout, key = connection(settings)
    salt = salt or os.urandom(16).hex()
    return {**resolved, "credential_salt": salt, "credential_digest": credential_digest(key, salt)}


def http_client(settings, **kwargs):
    import httpx

    resolved, headers, timeout, _key = connection(settings)
    return httpx.AsyncClient(
        base_url=resolved["endpoint"] + "/", headers=headers, timeout=timeout, **kwargs
    )


def private_json(path: Path, data, *, exclusive=False):
    """Write without exposing partial JSON or replacing a previous exercise."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
    if path.is_symlink():
        raise RuntimeError("Refusing a symlink for tutorial state")
    content = json.dumps(data, indent=2, allow_nan=False) + "\n"
    if exclusive:
        try:
            fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
        except FileExistsError as exc:
            raise RuntimeError("State already exists; inspect or recover the prior run") from exc
        with os.fdopen(fd, "w") as output:
            output.write(content)
            output.flush()
            os.fsync(output.fileno())
        return
    fd, temporary = tempfile.mkstemp(prefix=path.name + ".", dir=path.parent)
    try:
        with os.fdopen(fd, "w") as output:
            output.write(content)
            output.flush()
            os.fsync(output.fileno())
        os.replace(temporary, path)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


class TutorialState:
    def __init__(self, path, data):
        self.path = Path(path)
        self.data = data

    @classmethod
    def create(cls, path, settings, lesson, *, workspace_label=None, workspace_owner=None):
        data = {
            "schema_version": 1,
            "lesson": lesson,
            "identity": identity(settings),
            "run_token": uuid4().hex,
            "started_at": datetime.now(timezone.utc).isoformat(),
            "resources": {},
            "pending": None,
            # Written durably by begin() before the first remote mutation.
            # Missing in an older ledger means unknown, never proof of no writes.
            "remote_write_started": False,
            "operator_context": {
                "workspace_label": workspace_label,
                "workspace_owner": workspace_owner,
            },
        }
        private_json(Path(path), data, exclusive=True)
        return cls(path, data)

    @classmethod
    def load(cls, path, settings, lesson=None):
        path = Path(path)
        data = read_state(path)
        stored = data["identity"]
        salt = stored.get("credential_salt")
        if not isinstance(salt, str) or not re.fullmatch(r"[0-9a-f]{32}", salt):
            raise RuntimeError("Tutorial identity is missing its credential salt; do not reseed")
        if stored != identity(settings, salt):
            raise RuntimeError(
                "Endpoint, workspace, or credential changed; verify ownership before recovery"
            )
        if lesson is not None and data.get("lesson") != lesson:
            raise RuntimeError("State belongs to another tutorial")
        return cls(path, data)

    def save(self):
        private_json(self.path, self.data)

    def record(self, kind, resource_id, **metadata):
        if not isinstance(resource_id, str) or not resource_id.strip():
            raise RuntimeError(f"Missing returned {kind} ID; retained pending operation")
        resources = self.data["resources"].setdefault(kind, {})
        resources.setdefault(resource_id, {"status": "retained"}).update(metadata)
        self.save()

    def begin(self, operation):
        if self.data.get("pending"):
            raise RuntimeError("Unresolved write; inspect the retained state before retrying")
        self.data["remote_write_started"] = True
        self.data["pending"] = operation
        self.save()

    def finish(self):
        self.data["pending"] = None
        self.save()

    def inspect(self):
        data = json.loads(json.dumps(self.data))
        data["identity"].pop("credential_salt", None)
        data["identity"].pop("credential_digest", None)
        return data


def remember_message(state, message):
    role = getattr(message.role, "value", message.role)
    state.record("message", str(message.id), content_sha256=digest(message.content), role=role)


async def verify_messages(client, state):
    """Read-only exact-ID/content verification; safe to run in a fresh process."""
    expected = state.data["resources"].get("message", {})
    if not expected or not state.data.get("conversation_id"):
        raise RuntimeError("No complete message seed to verify")
    history = await client.short_term.get_conversation(state.data["conversation_id"])
    actual = {str(m.id): m for m in history.messages}
    if set(actual) != set(expected):
        raise RuntimeError("Conversation has missing or unrecorded messages; retained state")
    for message_id, entry in expected.items():
        message = actual[message_id]
        role = getattr(message.role, "value", message.role)
        if digest(message.content) != entry["content_sha256"] or role != entry["role"]:
            raise RuntimeError(f"Message readback differs: {message_id}")
    print(f"Verified: exact IDs, roles and content for {len(expected)} stored messages")


def main(argv=None):
    """Local diagnostics only; importing this module still performs no I/O."""
    import argparse

    parser = argparse.ArgumentParser(description="Inspect hosted lesson state without credentials")
    commands = parser.add_subparsers(dest="command", required=True)
    inspect = commands.add_parser("inspect-file", help="Print a redacted local summary; no network")
    inspect.add_argument("state", type=Path)
    retry = commands.add_parser("check-new-run", help="Check a separate path; preserve both files")
    retry.add_argument("state", type=Path)
    retry.add_argument("--next-state", type=Path, required=True)
    args = parser.parse_args(argv)
    if args.command == "inspect-file":
        print(json.dumps(inspect_file(args.state), indent=2))
    else:
        status = check_new_run(args.state, args.next_state)
        print(f"Recorded new-run check: {status}; original state preserved")
        print(f"Use --state {args.next_state} only after rechecking the lesson prerequisites")
        print("No service state was checked and no new file was created")


if __name__ == "__main__":
    main()

"""Scoped cleanup using the hosted API's supported conversation/entity routes."""

from __future__ import annotations

import asyncio
import time
from datetime import datetime, timezone
from urllib.parse import quote

# Start from retained message IDs, then inspect all sources and neighboring IDs.
PROVENANCE_QUERY = """
MATCH (m) WHERE m.id IN $message_ids
MATCH (m)-[:EXTRACTED_FROM]-(e:Entity)
WITH DISTINCT e
OPTIONAL MATCH (e)-[:EXTRACTED_FROM]-(source)
WITH e, collect(DISTINCT source.id) AS sources, count(DISTINCT source) AS source_count
OPTIONAL MATCH (e)--(neighbor)
RETURN e.id AS id, e.createdAt AS created_at, sources, source_count,
       collect(DISTINCT neighbor.id) AS neighbors, count(DISTINCT neighbor) AS neighbor_count
"""
RESIDUAL_QUERY = """
MATCH (n) WHERE n.id IN $ids OR n.conversationId = $conversation
RETURN n.id AS id, labels(n) AS labels
"""


async def request_json(http, method, route, **kwargs):
    response = await http.request(method, route, **kwargs)
    response.raise_for_status()
    return response.json()


async def query(http, cypher, params):
    result = await request_json(http, "POST", "query", json={"cypher": cypher, "params": params})
    rows = result.get("rows")
    if not isinstance(rows, list):
        raise RuntimeError("Query did not return rows; retained tutorial state")
    return rows


async def wait_for_extraction(http, conversation_id, count, *, timeout=60.0, interval=1.0):
    if timeout <= 0 or interval <= 0:
        raise ValueError("Polling timeout and interval must be positive")
    deadline = time.monotonic() + timeout
    route = f"conversations/{quote(conversation_id, safe='')}/extraction-status"
    while True:
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            raise TimeoutError("Extraction did not settle; no cleanup was attempted")
        result = await asyncio.wait_for(request_json(http, "GET", route), remaining)
        summary = result.get("summary")
        if not isinstance(summary, dict) or any(
            type(value) is not int or value < 0 for value in summary.values()
        ):
            raise RuntimeError("Invalid extraction status; no cleanup was attempted")
        if any(summary.get(key, 0) for key in ("failed", "error", "cancelled")):
            raise RuntimeError("Extraction failed; retain state for inspection")
        terminal = {"done", "completed", "skipped"}
        pending = {"pending", "processing", "running", "in_progress", "queued"}
        if set(summary) - terminal - pending:
            raise RuntimeError("Unknown extraction status; retain state for inspection")
        if sum(summary.values()) == count and not any(summary.get(key, 0) for key in pending):
            return
        await asyncio.sleep(min(interval, max(0, deadline - time.monotonic())))


async def delete_and_verify(http, route):
    response = await http.delete(route)
    if response.status_code != 404:
        response.raise_for_status()
    readback = await http.get(route)
    if readback.status_code != 404:
        readback.raise_for_status()
        raise RuntimeError(f"Resource still exists after DELETE: {route}")


def exclusively_owned(row, message_ids, started_at):
    try:
        created = datetime.fromisoformat(row["created_at"].replace("Z", "+00:00"))
        started = datetime.fromisoformat(started_at)
        # collect(node.id) omits anonymous nodes. Compare node counts as well
        # so missing IDs cannot hide another source or adjacent resource.
        for field, count_field in (("sources", "source_count"), ("neighbors", "neighbor_count")):
            ids = row.get(field)
            count = row.get(count_field)
            if (
                not isinstance(ids, list)
                or any(not isinstance(value, str) or not value for value in ids)
                or type(count) is not int
                or count != len(set(ids))
            ):
                return False
        return (
            bool(row.get("sources"))
            and set(row["sources"]) <= message_ids
            and started <= created <= datetime.now(timezone.utc)
        )
    except (KeyError, ValueError, TypeError):
        return False


async def cleanup(http, state, *, timeout=60.0, interval=1.0) -> bool:
    """Retain shared/uncertain records; never use unsupported step/skill deletion."""
    if state.data.get("pending"):
        raise RuntimeError("Unresolved write; reconcile its outcome before cleanup")
    conversation_id = state.data.get("conversation_id")
    resources = state.data["resources"]
    if not conversation_id or conversation_id not in resources.get("conversation", {}):
        raise RuntimeError("No recorded conversation owned by this run")
    # A previous preflight cannot authorize a later deletion. Ownership and
    # neighboring sources may have changed while a failed run was interrupted.
    for entry in resources.get("entity", {}).values():
        if entry.get("status") != "deleted":
            entry["owned"] = False
    state.data["cleanup_complete"] = False
    state.save()
    route = f"conversations/{quote(conversation_id, safe='')}"
    response = await http.get(route)
    owned_messages = set(resources.get("message", {}))
    conversation = resources["conversation"][conversation_id]
    if response.status_code != 404:
        response.raise_for_status()
        metadata = response.json()
        expected_metadata = conversation.get("metadata", {})
        if (
            expected_metadata.get("tutorialRun") != state.data["run_token"]
            or metadata.get("id") != conversation_id
            or metadata.get("metadata") != expected_metadata
        ):
            raise RuntimeError("Conversation ownership mismatch; no cleanup attempted")
        history = await request_json(http, "GET", route + "/messages")
        messages = history.get("messages") if isinstance(history, dict) else history
        if not isinstance(messages, list) or {m.get("id") for m in messages} != owned_messages:
            raise RuntimeError("Unrecorded or missing messages; no cleanup attempted")
        await wait_for_extraction(
            http, conversation_id, len(owned_messages), timeout=timeout, interval=interval
        )
        rows = await query(http, PROVENANCE_QUERY, {"message_ids": sorted(owned_messages)})
        candidates = {
            r["id"] for r in rows if exclusively_owned(r, owned_messages, state.data["started_at"])
        }
        while True:
            exclusive = {
                r["id"]
                for r in rows
                if r["id"] in candidates
                and isinstance(r.get("neighbors"), list)
                and set(r["neighbors"]) <= owned_messages | candidates
            }
            if exclusive == candidates:
                break
            candidates = exclusive
        for row in rows:
            # A graph edge to data outside this run prevents entity deletion too.
            owned = row["id"] in candidates
            state.record("entity", row["id"], status="retained", owned=owned, provenance=row)
        state.data["cleanup_preflight_complete"] = True
        state.save()
    elif not state.data.get("cleanup_preflight_complete"):
        raise RuntimeError(
            "Conversation disappeared before provenance capture; inspect retained IDs"
        )

    for entity_id, entry in resources.get("entity", {}).items():
        entity_route = "entities/" + quote(entity_id, safe="")
        if entry.get("status") == "deleted" or not entry.get("owned"):
            readback = await http.get(entity_route)
            if readback.status_code == 404:
                entry["status"] = "deleted"
                entry["disposition"] = "absence independently verified"
                state.save()
                continue
            readback.raise_for_status()
            entry["status"] = "retained"
            entry["owned"] = False
        if not entry.get("owned"):
            entry["disposition"] = "retained: shared or insufficient ownership evidence"
            state.save()
            continue
        await delete_and_verify(http, entity_route)
        entry["status"] = "deleted"
        state.save()
    await delete_and_verify(http, route)
    conversation["status"] = "deleted"
    state.save()
    ids = [
        conversation_id,
        *owned_messages,
        *(
            key
            for key, value in resources.get("entity", {}).items()
            if value["status"] == "deleted"
        ),
        *(
            key
            for kind, entries in resources.items()
            if kind not in {"conversation", "message", "entity"}
            for key in entries
        ),
    ]
    rows = await query(http, RESIDUAL_QUERY, {"ids": ids, "conversation": conversation_id})
    state.data["remaining_rows"] = rows
    remaining_ids = {row.get("id") for row in rows}
    for entries in resources.values():
        for resource_id, entry in entries.items():
            if resource_id in remaining_ids:
                entry["status"] = "retained"
                entry["disposition"] = "present in final exact-ID readback"
    for message_id, entry in resources.get("message", {}).items():
        if message_id not in remaining_ids:
            entry["status"] = "deleted"
    retained = {
        kind: [key for key, value in entries.items() if value["status"] != "deleted"]
        for kind, entries in resources.items()
    }
    state.data["retained_resources"] = {kind: ids for kind, ids in retained.items() if ids}
    state.save()
    # Steps/skills can remain attached to the deleted conversation. These are
    # recorded dispositions, not a successful assertion that every row vanished.
    unsupported = {
        key
        for kind, entries in resources.items()
        if kind not in {"conversation", "message", "entity"}
        for key in entries
    }
    if any(row.get("id") not in unsupported for row in rows):
        raise RuntimeError("Unexpected run-owned residuals remain; inspect state")
    state.data["cleanup_complete"] = not state.data["retained_resources"]
    state.save()
    print("Verified: owned conversation and proven-owned entities are absent")
    print(f"Retained resources: {state.data['retained_resources']}")
    return not state.data["retained_resources"]