#!/usr/bin/env python3
"""muse_session_bind.py — Per-session credential isolation (P3).

Each muse session gets its own config root::

    /tmp/muse-session-<pid>/muse/
      <everything symlinked from the global config EXCEPT auth.json>
      auth.json            <- COPY of the bound profile's credentials (0600)
    /tmp/muse-session-<pid>/bind.json  <- {profile, pid, created, auth_src}

``launch`` execs muse with XDG_CONFIG_HOME pointed at the session dir
(the binary resolves its config root as $XDG_CONFIG_HOME/muse, else
$HOME/.config/muse), so switching profiles never disturbs live
sessions: the fleet-wide 400 outage class disappears by construction.
Exec (not supervise) preserves the pane's ``muse-bin`` identity, so
watcher coverage and ``box runtime`` keep working unchanged.

Token refreshes land in the session copy. ``save``/``reap`` copy newer
bytes back to the profile store (newest-wins across concurrent
sessions sharing a profile; nothing is ever written to the legacy
global auth.json). Dead sessions are reaped by scan, so kill -9 loses
nothing but promptness.

Companion to the peer's muse_resume_pool (which reads
session_profiles.json): ``launch --session-id`` records the binding
there for future resume guards.
"""

import argparse
import hashlib
import json
import os
import shutil
import sys
import time
from datetime import datetime, timezone

SESSION_PREFIX = "muse-session-"
BIND_FILENAME = "bind.json"
AUTH_FILENAME = "auth.json"


def default_config_src():
    """Global config source (explicit env wins, else the real home)."""
    return (os.environ.get("MUSE_CONFIG_SRC")
            or os.path.join(os.path.expanduser("~"), ".config", "muse"))


def session_dir_for(parent, pid):
    return os.path.join(parent, "%s%d" % (SESSION_PREFIX, pid))


def _now():
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def _pid_alive(pid):
    try:
        os.kill(pid, 0)
        return True
    except Exception:
        return False


def _pid_is_muse(pid):
    """True if pid's cmdline looks like a muse session (pid-reuse guard)."""
    try:
        with open("/proc/%d/cmdline" % pid, "rb") as f:
            cmd = f.read().decode(errors="replace").lower()
        return "muse-bin" in cmd or "muse-code" in cmd
    except Exception:
        return False


def _fingerprint(path):
    """Short sha256 of a credential file for logs (never the bytes)."""
    try:
        h = hashlib.sha256()
        with open(path, "rb") as f:
            h.update(f.read())
        return h.hexdigest()[:12]
    except OSError:
        return "missing"


def _write_private_bytes(path, data):
    """Write bytes with 0600 perms, atomically. Returns True on success."""
    try:
        tmp = "%s.tmp.%d" % (path, os.getpid())
        fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
        try:
            os.write(fd, data)
            os.fsync(fd)
        finally:
            os.close(fd)
        os.replace(tmp, path)
        return True
    except OSError:
        return False


def profile_auth_path(config_src, profile):
    return os.path.join(config_src, "accounts", profile, AUTH_FILENAME)


def read_bind(sessdir):
    try:
        with open(os.path.join(sessdir, BIND_FILENAME)) as f:
            data = json.load(f)
        return data if isinstance(data, dict) else None
    except (OSError, ValueError):
        return None


def session_liveness(sessdir):
    """live | dead | unknown (no/invalid bind record: never reap)."""
    bind = read_bind(sessdir)
    if not bind or not isinstance(bind.get("pid"), int):
        return "unknown"
    pid = bind["pid"]
    if _pid_alive(pid) and _pid_is_muse(pid):
        return "live"
    return "dead"


def list_bound(parent="/tmp"):
    """Session dirs carrying our bind record (foreign dirs ignored)."""
    out = []
    try:
        names = sorted(os.listdir(parent))
    except OSError:
        return out
    for name in names:
        if not name.startswith(SESSION_PREFIX):
            continue
        sessdir = os.path.join(parent, name)
        if not os.path.isdir(sessdir):
            continue
        if read_bind(sessdir) is None:
            continue
        out.append(sessdir)
    return out


def build_session_dir(parent, pid, config_src, auth_src, profile):
    """Create the isolated config root. Returns sessdir.

    Raises RuntimeError when the slot is held by a live session, or
    OSError/ValueError for missing sources.
    """
    auth_src = os.path.realpath(auth_src)
    if not os.path.isfile(auth_src):
        raise ValueError("no credentials at %s" % auth_src)
    if not os.path.isdir(config_src):
        raise ValueError("no config source at %s" % config_src)
    sessdir = session_dir_for(parent, pid)
    cfgdir = os.path.join(sessdir, "muse")
    if os.path.exists(sessdir):
        if session_liveness(sessdir) == "live":
            raise RuntimeError("session slot %s is live" % sessdir)
        shutil.rmtree(sessdir, ignore_errors=True)
    os.makedirs(cfgdir)
    for entry in sorted(os.listdir(config_src)):
        if entry == AUTH_FILENAME:
            continue
        target = os.path.join(config_src, entry)
        try:
            os.symlink(target, os.path.join(cfgdir, entry))
        except OSError:
            pass
    with open(auth_src, "rb") as f:
        creds = f.read()
    if not _write_private_bytes(os.path.join(cfgdir, AUTH_FILENAME), creds):
        raise OSError("cannot plant auth.json in %s" % cfgdir)
    with open(os.path.join(sessdir, BIND_FILENAME), "w") as f:
        json.dump({"profile": profile, "pid": pid,
                   "created": _now(), "auth_src": auth_src}, f, indent=1)
    return sessdir


def record_session_profile(config_src, session_id, profile):
    """Note session->profile for resume guards. Returns True on success."""
    path = os.path.join(config_src, "session_profiles.json")
    try:
        with open(path) as f:
            data = json.load(f)
        if not isinstance(data, dict):
            data = {}
    except (OSError, ValueError):
        data = {}
    data[str(session_id)] = str(profile)
    try:
        tmp = "%s.tmp.%d" % (path, os.getpid())
        with open(tmp, "w") as f:
            json.dump(data, f, indent=1)
        os.replace(tmp, path)
        return True
    except OSError:
        return False


def save_session(sessdir, config_src=None):
    """Sync a session copy back to its profile when newer.

    Returns {"status", ...}; statuses: synced | skipped-stale |
    skipped-missing | no-bind. Never raises, never logs token bytes.
    """
    config_src = config_src or default_config_src()
    bind = read_bind(sessdir)
    if not bind or not bind.get("profile"):
        return {"status": "no-bind", "sessdir": sessdir}
    profile = bind["profile"]
    sess_auth = os.path.join(sessdir, "muse", AUTH_FILENAME)
    dest = os.path.realpath(profile_auth_path(config_src, profile))
    if not os.path.isfile(sess_auth):
        return {"status": "skipped-missing", "sessdir": sessdir,
                "profile": profile}
    try:
        sess_mtime = os.path.getmtime(sess_auth)
    except OSError:
        return {"status": "skipped-missing", "sessdir": sessdir,
                "profile": profile}
    try:
        dest_mtime = os.path.getmtime(dest)
    except OSError:
        dest_mtime = -1
    if dest_mtime >= sess_mtime:
        return {"status": "skipped-stale", "sessdir": sessdir,
                "profile": profile, "session_fp": _fingerprint(sess_auth),
                "profile_fp": _fingerprint(dest)}
    try:
        with open(sess_auth, "rb") as f:
            creds = f.read()
    except OSError:
        return {"status": "skipped-missing", "sessdir": sessdir,
                "profile": profile}
    try:
        os.makedirs(os.path.dirname(dest), exist_ok=True)
    except OSError:
        pass
    if not _write_private_bytes(dest, creds):
        return {"status": "error", "sessdir": sessdir, "profile": profile}
    return {"status": "synced", "sessdir": sessdir, "profile": profile,
            "session_fp": _fingerprint(sess_auth),
            "profile_fp": _fingerprint(dest)}


def reap(parent="/tmp", config_src=None):
    """Sync + remove dead bound sessions. Returns {"reaped", "live"}."""
    config_src = config_src or default_config_src()
    reaped, live = [], []
    for sessdir in list_bound(parent):
        if session_liveness(sessdir) == "live":
            live.append(sessdir)
            continue
        res = save_session(sessdir, config_src)
        shutil.rmtree(sessdir, ignore_errors=True)
        reaped.append({"sessdir": sessdir, "save": res["status"],
                       "profile": res.get("profile")})
    return {"reaped": reaped, "live": live}


def launch(profile=None, auth_file=None, session_id=None, config_src=None,
           cmd=None, parent="/tmp", dry_run=False, _exec=os.execvpe):
    """Bind then exec. With dry_run, return the plan without exec'ing."""
    config_src = config_src or default_config_src()
    if auth_file:
        auth_src = os.path.realpath(auth_file)
    elif profile:
        auth_src = profile_auth_path(config_src, profile)
    else:
        raise ValueError("need --profile or --auth-file")
    if not os.path.isfile(auth_src):
        raise ValueError("no credentials at %s" % auth_src)
    pid = os.getpid()
    if dry_run:
        return {"sessdir": session_dir_for(parent, pid),
                "xdg_config_home": session_dir_for(parent, pid),
                "profile": profile, "auth_src": auth_src,
                "cmd": cmd or []}
    # Reap BEFORE building: our own fresh dir would read as dead (the
    # launcher is python, not muse, until it execs) and eat itself.
    reap(parent=parent, config_src=config_src)
    sessdir = build_session_dir(parent, pid, config_src, auth_src,
                                profile or "explicit")
    if session_id:
        record_session_profile(config_src, session_id,
                               profile or "explicit")
    env = dict(os.environ)
    env["XDG_CONFIG_HOME"] = sessdir
    env["MUSE_SESSION_BIND_DIR"] = sessdir
    _exec(cmd[0], cmd, env)
    return None  # unreachable; exec replaces the image


def cmd_status(args):
    config_src = args.config_src or default_config_src()
    rows = []
    for sessdir in list_bound(args.parent):
        bind = read_bind(sessdir) or {}
        sess_auth = os.path.join(sessdir, "muse", AUTH_FILENAME)
        prof_auth = profile_auth_path(config_src, bind.get("profile", ""))
        rows.append({"sessdir": sessdir, "profile": bind.get("profile"),
                     "pid": bind.get("pid"),
                     "liveness": session_liveness(sessdir),
                     "session_fp": _fingerprint(sess_auth),
                     "profile_fp": _fingerprint(prof_auth)})
    if args.json:
        print(json.dumps({"sessions": rows}, indent=1))
    else:
        if not rows:
            print("No bound sessions under %s." % args.parent)
            return 0
        for r in rows:
            print("%s profile=%s pid=%s %s session=%s profile=%s" % (
                r["sessdir"], r["profile"], r["pid"], r["liveness"],
                r["session_fp"], r["profile_fp"]))
    return 0


def main(argv=None):
    ap = argparse.ArgumentParser(
        prog="muse_session_bind",
        description="Per-session credential isolation for muse.")
    ap.add_argument("--parent", default="/tmp",
                    help="Session dir parent (default /tmp).")
    ap.add_argument("--config-src", default=None,
                    help="Global config source (default ~/.config/muse).")
    sub = ap.add_subparsers(dest="command", required=True)

    p_l = sub.add_parser("launch", help="Bind a profile, then exec muse.")
    p_l.add_argument("--profile", default=None)
    p_l.add_argument("--auth-file", default=None)
    p_l.add_argument("--session-id", default=None)
    p_l.add_argument("--dry-run", action="store_true")
    p_l.add_argument("cmd", nargs=argparse.REMAINDER,
                     help="Command after --, e.g. -- muse-code")

    p_s = sub.add_parser("save", help="Sync session tokens back to profile.")
    g = p_s.add_mutually_exclusive_group(required=True)
    g.add_argument("--pid", type=int)
    g.add_argument("--dir")
    g.add_argument("--all", action="store_true")
    p_s.add_argument("--json", action="store_true")

    p_r = sub.add_parser("reap", help="Sync + remove dead sessions.")
    p_r.add_argument("--json", action="store_true")

    p_st = sub.add_parser("status", help="List bound sessions.")
    p_st.add_argument("--json", action="store_true")

    args = ap.parse_args(argv)
    config_src = args.config_src or default_config_src()
    if args.command == "launch":
        cmd = [c for c in args.cmd if c != "--"]
        if not cmd and not args.dry_run:
            print("launch needs a command: launch ... -- muse-code [...]",
                  file=sys.stderr)
            return 2
        try:
            plan = launch(profile=args.profile, auth_file=args.auth_file,
                          session_id=args.session_id, config_src=config_src,
                          cmd=cmd, parent=args.parent,
                          dry_run=args.dry_run)
        except (ValueError, RuntimeError, OSError) as e:
            print("launch refused: %s" % (e,), file=sys.stderr)
            return 1
        if args.dry_run:
            print(json.dumps(plan, indent=1))
        return 0
    if args.command == "save":
        if args.pid is not None:
            targets = [session_dir_for(args.parent, args.pid)]
        elif args.dir:
            targets = [args.dir]
        else:
            targets = list_bound(args.parent)
        results = [save_session(t, config_src) for t in targets]
        if args.json:
            print(json.dumps({"saved": results}, indent=1))
        else:
            for r in results:
                print("%s: %s" % (r["sessdir"], r["status"]))
        return 0
    if args.command == "reap":
        res = reap(parent=args.parent, config_src=config_src)
        if args.json:
            print(json.dumps(res, indent=1))
        else:
            for r in res["reaped"]:
                print("reaped %s (%s)" % (r["sessdir"], r["save"]))
            if not res["reaped"]:
                print("Nothing to reap.")
        return 0
    if args.command == "status":
        return cmd_status(args)
    return 2


if __name__ == "__main__":
    sys.exit(main())