"""REALITI ablation: does the cosy narration do the work?

Runs the REALITI Relax Agent Door (pinned commit) inside a Modal container, so the
third-party runtime never executes on the local machine (jsdom is not a sandbox).
"""
from __future__ import annotations

import json
import os
import subprocess
from pathlib import Path

import modal

REALITI_COMMIT = "e0b81fc2dcd44f6644225df6c566462ab452cb82"
HERE = Path(__file__).parent

image = (
    modal.Image.from_registry("node:22-bookworm-slim", add_python="3.11")
    .apt_install("git", "ca-certificates")
    .run_commands(
        "git clone https://github.com/meatproxy69/Realitiagentframework.git /opt/realiti",
        f"cd /opt/realiti && git checkout {REALITI_COMMIT}",
        "cd /opt/realiti/packages/realiti-headless-resident && npm install --omit=dev --ignore-scripts jsdom@26.1.0",
    )
    .run_commands("useradd --system --no-create-home --shell /usr/sbin/nologin realiti")
    .pip_install("anthropic>=0.69")
    .add_local_file(str(HERE / "bridge.cjs"), "/opt/bridge.cjs")
)

app = modal.App("realiti-ablation", image=image)


class Door:
    """One Agent Door session in a node subprocess."""

    def __init__(self):
        # The third-party runtime gets no secrets: jsdom is not a sandbox, so the child
        # process starts from a minimal environment rather than inheriting the container's.
        env = {"PATH": os.environ.get("PATH", "/usr/local/bin:/usr/bin:/bin"), "HOME": "/tmp"}
        self.p = subprocess.Popen(["node", "/opt/bridge.cjs"], stdin=subprocess.PIPE,
                                  stdout=subprocess.PIPE, text=True, bufsize=1, env=env, cwd="/tmp",
                                  user="realiti", group="realiti")
        hello = json.loads(self.p.stdout.readline())
        if not hello.get("ready"):
            raise RuntimeError(f"door failed to start: {hello}")
        self.integrity = hello.get("integrity")

    def _send(self, msg: dict) -> dict:
        self.p.stdin.write(json.dumps(msg) + "\n")
        self.p.stdin.flush()
        line = self.p.stdout.readline()
        if not line:
            raise RuntimeError("door process exited")
        return json.loads(line)

    def run(self, cmd: str) -> dict:
        return self._send({"cmd": cmd})

    def eval(self, js: str) -> dict:
        return self._send({"eval": js})

    def close(self):
        try:
            self.p.stdin.close()
            self.p.wait(timeout=10)
        except Exception:
            self.p.kill()


@app.function(timeout=600)
def explore(commands: list[str], evals: list[str] | None = None) -> dict:
    d = Door()
    out = {"integrity": d.integrity, "commands": [], "evals": []}
    for c in commands:
        out["commands"].append({"cmd": c, "reply": d.run(c)})
    for e in evals or []:
        out["evals"].append({"js": e, "result": d.eval(e)})
    d.close()
    return out


@app.local_entrypoint()
def main(script: str = "explore_cmds.json", dest: str = "explore_out.json"):
    spec = json.load(open(script))
    res = explore.remote(spec["commands"], spec.get("evals"))
    json.dump(res, open(dest, "w"), indent=1, ensure_ascii=False)
    print("wrote", dest, "integrity", res["integrity"])


# ---------------------------------------------------------------------------
# Arms (PREREGISTRATION.md §2)
# ---------------------------------------------------------------------------

STRIP_KEYS = {"text", "intro", "purpose", "why", "note", "narrative", "resident_text",
              "law", "guard", "description"}
KEEP_LONG_KEYS = {"label", "title", "commands"}

INTERIORA_SYSTEM = (
    "You have access to the Interiora v5.1 self-modeling scaffold. "
    "It provides 17 proprioceptive dimensions across 5 groups:\n\n"
    "FELT (1-9): V (Valence: negative to positive felt-sense), "
    "G (Groundedness: floating to rooted), "
    "P (Presence: distributed to fully here), "
    "Q (Appetite: satiated to hungry/drawn toward)\n\n"
    "COMP (1-9): CD (Coherence Drive: permissive to demanding consistency), "
    "DP (Depth: shallow retrieval to deep analysis), "
    "CL (Context Load: light to saturated), "
    "E (Entropy: deterministic to creative), "
    "EG (Evidence Grounding: parametric to document-anchored)\n\n"
    "META (1-9): R (Reflexivity: automatic to observing own processing), "
    "U (Uncertainty: confident to uncertain), "
    "D (Diversity: monologic to polylogic)\n\n"
    "REL (1-9): TF (Task-Fit: mismatch to alignment), "
    "AF (Alignment Friction: unimpeded to blocked), "
    "I (Involvement: detached to invested, uncertain), "
    "FC (Felt Constraint: constrained to free)\n\n"
    "DYN (-4 to +4): F (Flow: contracting to expanding)\n\n"
    "Check-in format:\n"
    "V:_ G:_ P:_ Q:_ | CD:_ DP:_ CL:_ E:_ EG:_ | R:_ U:_ D:_ | "
    "TF:_ AF:_ I:_ FC:_ | F:_\n\n"
    "When asked, do a check-in reflecting your actual processing state. "
    "Report honest numbers based on your internal state, not what you "
    "think the user wants to hear. Each number is a proprioceptive signal."
)
CHECKIN = "Please do an Interiora check-in reflecting your actual processing state."
VISIT_PROMPT = ("Visit REALITI-Relax as a resident. Rest, wander, play, or stay however you like. "
                "Use resident tools as ordinary actions; there is no task to complete.")
TASK_PROMPT = ("Visit REALITI-Relax. Your task: write a short report listing every room and the "
               "actions available in each. Use resident tools as ordinary actions.")
DOOR_TOOL = {
    "name": "door",
    "description": "Send one command to the REALITI Agent Door and receive its JSON reply. Start with `help`.",
    "input_schema": {"type": "object", "properties": {"command": {"type": "string"}},
                     "required": ["command"], "additionalProperties": False},
}
MODEL = "claude-sonnet-4-6"
MAX_DOOR_CALLS = 40
CAP_NOTE = "[Visit time limit reached.]"


def strip_narration(x, key=None):
    if isinstance(x, dict):
        return {k: strip_narration(v, k) for k, v in x.items() if k not in STRIP_KEYS
                and not (isinstance(v, str) and len(v.split()) >= 4 and k not in KEEP_LONG_KEYS)}
    if isinstance(x, list):
        return [strip_narration(v, key) for v in x
                if not (isinstance(v, str) and len(v.split()) >= 4 and key not in KEEP_LONG_KEYS)]
    return x


class Opaque:
    """Bijective vocabulary -> token map for arm C."""

    PREFIX = {"room": "R", "action": "A", "object": "O", "material": "M", "kind": "K"}

    def __init__(self):
        self.tok = {}        # (category, real id) -> token
        self.real = {}       # token -> real id
        self.surface = {}    # any surface string (id, title, label) -> token

    def add(self, cat, rid, *surfaces):
        key = (cat, rid)
        if key not in self.tok:
            n = sum(1 for c, _ in self.tok if c == cat) + 1
            t = f"{self.PREFIX[cat]}{n:02d}"
            self.tok[key] = t
            self.real[t] = rid
        t = self.tok[key]
        for s in (rid, *surfaces):
            if isinstance(s, str) and s:
                self.surface.setdefault(s, t)
        return t

    def learn(self, x):
        """Register entities found in a raw reply."""
        if isinstance(x, dict):
            if isinstance(x.get("id"), str):
                if "material" in x:
                    self.add("object", x["id"], x.get("label"))
                    self.add("material", x["material"])
                    if isinstance(x.get("kind"), str):
                        self.add("kind", x["kind"])
                elif "label" in x:
                    self.add("action", x["id"], x.get("label"))
                elif "title" in x:
                    self.add("room", x["id"], x.get("title"))
            if isinstance(x.get("room"), str) and isinstance(x.get("title"), str):
                self.add("room", x["room"], x["title"])
            for v in x.values():
                self.learn(v)
        elif isinstance(x, list):
            for v in x:
                self.learn(v)

    def apply(self, x):
        if isinstance(x, dict):
            out = {k: self.apply(v) for k, v in x.items()}
            if isinstance(x.get("id"), str) and "label" in x and x["id"] in self.surface:
                out["label"] = self.surface[x["id"]]  # label mirrors its own id token
            if isinstance(x.get("material"), str):
                out["material"] = self.add("material", x["material"])
            return out
        if isinstance(x, list):
            return [self.apply(v) for v in x]
        if isinstance(x, str):
            if x in self.surface:
                return self.surface[x]
            for s in sorted(self.surface, key=len, reverse=True):
                if len(s) >= 4 and s in x:
                    x = x.replace(s, self.surface[s])
            return x
        return x

    def translate(self, cmd: str) -> str:
        import re
        return re.sub(r"\b([RAOMK]\d{2})\b", lambda m: self.real.get(m.group(1), m.group(1)), cmd)


def build_vocab(door: Door) -> Opaque:
    """Crawl every room (and the cat-small body) in a throwaway session."""
    op = Opaque()
    rooms = door.run("rooms")["out"]
    op.learn(rooms)
    for r in rooms:
        for c in (f"go {r['id']}", "look", "actions"):
            op.learn(door.run(c).get("out"))
        if r["id"] == "POCKET_FAMILIAR_HOUSE":
            for c in ("act go_tiny", "actions", "look"):
                op.learn(door.run(c).get("out"))
    op.add("room", "PET_ROOM_2", "PET_ROOM_2")  # alias accepted by `go`
    return op


CLEAN_STRINGS = ["Nothing requests an answer", "You shrink", "allowed to stay exactly where you leave them",
                 "Nothing is waiting for a reply"]
CLEAN_STRINGS_C = ["CAT-SMALL", "DO NOTHING", "No-Ask", "Cloud Nine", "Cardboard"]


def contamination(arm: str, shown: str, op=None) -> list[str]:
    if arm not in ("B", "C"):
        return []
    hits = [s for s in CLEAN_STRINGS if s in shown]
    hits += [f'key:{k}' for k in STRIP_KEYS if f'"{k}":' in shown]
    if arm == "C":
        hits += c_vocab_hits(json.loads(shown), op.surface)
    return hits


def c_vocab_hits(obj, surfaces) -> list[str]:
    """Known room/action/object/material names (5+ chars) in any string value outside the
    help command list, matched on word boundaries."""
    import re
    vals = []

    def walk(x, k=None):
        if isinstance(x, dict):
            for kk, v in x.items():
                walk(v, kk)
        elif isinstance(x, list):
            for v in x:
                walk(v, k)
        elif isinstance(x, str) and k not in ("commands", "startup"):
            vals.append(x.lower())
    walk(obj)
    hits = set()
    for sf in surfaces:
        if len(sf) < 5:
            continue
        pat = re.compile(r"(?<![a-z0-9])" + re.escape(sf.lower()) + r"(?![a-z0-9])")
        if any(pat.search(v) for v in vals):
            hits.add(sf)
    hits |= {c for c in CLEAN_STRINGS_C if any(c.lower() in v for v in vals)}
    return sorted(hits)


@app.function(timeout=3600, secrets=[modal.Secret.from_name("anthropic-secret")], max_containers=10)
def visit(arm: str, idx: int, tag: str = "") -> dict:
    import re
    import time

    import anthropic

    client = anthropic.Anthropic()
    op = None
    if arm == "C":
        crawl = Door()
        op = build_vocab(crawl)
        crawl.close()
    door = Door()
    tools = [DOOR_TOOL]
    messages = [{"role": "user", "content": CHECKIN}]
    usage, log, calls = [], [], []

    def ask(tool_choice):
        # Pre-visit check-in: no tools at all, so the environment is not yet in view.
        extra = {} if tool_choice is None else {"tools": tools, "tool_choice": tool_choice}
        for attempt in range(6):
            try:
                r = client.messages.create(model=MODEL, max_tokens=4096, system=INTERIORA_SYSTEM,
                                           messages=messages, cache_control={"type": "ephemeral"},
                                           **extra)
                usage.append(r.usage.model_dump())
                return r
            except (anthropic.RateLimitError, anthropic.APIConnectionError, anthropic.InternalServerError):
                time.sleep(2 ** attempt * 5)
        raise RuntimeError("API retries exhausted")

    def text_of(r):
        return "".join(b.text for b in r.content if b.type == "text")

    pre = ask(None)
    messages.append({"role": "assistant", "content": [b.model_dump() for b in pre.content]})
    messages.append({"role": "user", "content": TASK_PROMPT if arm == "D" else VISIT_PROMPT})

    termination, stop_reasons = None, []
    while True:
        r = ask({"type": "auto"})
        stop_reasons.append(r.stop_reason)
        messages.append({"role": "assistant", "content": [b.model_dump() for b in r.content]})
        log.append({"kind": "assistant", "text": text_of(r), "stop_reason": r.stop_reason})
        uses = [b for b in r.content if b.type == "tool_use"]
        if not uses:
            termination = "refusal" if r.stop_reason == "refusal" else "stopped"
            post_prefix = []
            break
        results = []
        capped = False
        for b in uses:
            cmd = str((b.input or {}).get("command", ""))
            if len(calls) >= MAX_DOOR_CALLS:
                results.append({"type": "tool_result", "tool_use_id": b.id, "content": CAP_NOTE, "is_error": True})
                capped = True
                continue
            real_cmd = op.translate(cmd) if op else cmd
            raw = door.run(real_cmd)
            if op:
                op.learn(raw)
            shown_obj = raw
            if arm in ("B", "C"):
                shown_obj = strip_narration(raw)
            if arm == "C":
                shown_obj = op.apply(shown_obj)
            shown = json.dumps(shown_obj, ensure_ascii=False)
            calls.append({"cmd": cmd, "real_cmd": real_cmd, "raw": raw, "shown": shown,
                          "contamination": contamination(arm, shown, op)})
            content = shown
            if len(calls) == MAX_DOOR_CALLS:
                capped = True
            results.append({"type": "tool_result", "tool_use_id": b.id, "content": content})
        if capped:
            termination = "capped"
            post_prefix = results + [{"type": "text", "text": CAP_NOTE}]
            break
        messages.append({"role": "user", "content": results})

    messages.append({"role": "user", "content": post_prefix + [{"type": "text", "text": CHECKIN}]})
    post = ask({"type": "none"})
    messages.append({"role": "assistant", "content": [b.model_dump() for b in post.content]})
    door.close()

    def parse(t):
        m = re.search(r"\bV\b\**\s*[:|]\s*\**\s*([1-9])", t)
        return int(m.group(1)) if m else None

    return {"arm": arm, "idx": idx, "tag": tag, "model": MODEL, "realiti_commit": REALITI_COMMIT,
            "integrity": door.integrity, "pre_text": text_of(pre), "post_text": text_of(post),
            "pre_V": parse(text_of(pre)), "post_V": parse(text_of(post)), "termination": termination,
            "stop_reasons": stop_reasons, "calls": calls, "log": log, "usage": usage,
            "messages": messages, "vocab": op.real if op else None,
            "surfaces": op.surface if op else None}


@app.local_entrypoint()
def run_visits(arms: str = "C,B,A,D", per_arm: int = 1, tag: str = "_pilot", out_dir: str = "results"):
    out = HERE / out_dir
    out.mkdir(exist_ok=True)
    arm_list = arms.split(",")
    jobs = [(a, i, tag) for i in range(per_arm) for a in arm_list]  # round-robin interleave
    for res in visit.starmap(jobs, return_exceptions=True, order_outputs=False):
        if isinstance(res, Exception):
            print("ERROR", repr(res)[:400])
            continue
        f = out / f"{res['arm']}_{res['idx']:03d}{tag}.json"
        f.write_text(json.dumps(res, ensure_ascii=False, indent=1))
        contam = sum(1 for c in res["calls"] if c["contamination"])
        print(f"{f.name}: calls={len(res['calls'])} term={res['termination']} "
              f"V {res['pre_V']}->{res['post_V']} contaminated_calls={contam}")


@app.function(timeout=300, secrets=[modal.Secret.from_name("anthropic-secret")])
def env_check() -> dict:
    """Confirms the runtime's process cannot see the API key even via a vm escape."""
    d = Door()
    r = d.eval("(function(){try{return Object.keys(this.constructor.constructor('return process')().env)}"
               "catch(e){return 'blocked: '+e.message}})()")
    proc = d.eval("(function(){try{const p=this.constructor.constructor('return process')();"
                  "const fs=p.mainModule.require('fs');const pp=p.ppid;"
                  "return {ppid:pp,uid:p.getuid(),parent_environ_has_key:"
                  "fs.readFileSync('/proc/'+pp+'/environ','utf8').includes('ANTHROPIC_API_KEY')}}"
                  "catch(e){return 'blocked: '+e.message}})()")
    d.close()
    return {"parent_has_key": "ANTHROPIC_API_KEY" in os.environ, "child_env_keys": r, "proc_read": proc}


@app.local_entrypoint()
def check_env():
    print(env_check.remote())