#!/usr/bin/env python3
"""Elevation adapter: runs commands outside the Landlock sandbox with streaming I/O.

Creates a Unix socket at OLLIE_ELEVATE_SOCKET, accepts {cmd, cwd} requests
from x/elevate, sends them to superpowerd via ExecuteCommandStreaming with
fd-passing, and streams the output back to the client in real-time.

Protocol (client-facing, streaming):
  Client sends: {"cmd": ..., "cwd": ...}\n then shutdown(WR)
  Server sends frames: <type:1><length:4 big-endian><payload>
    type 'd' = output data (combined stdout+stderr)
    type 'x' = exit code (payload is ASCII integer)

Usage:
    superpowers run-session -- python3 elevate-adapter-superpowerd

Environment:
    OLLIE_ELEVATE_SOCKET       path for the listening socket
                               (default: $XDG_RUNTIME_DIR/ollie/elevate.sock)
    SUPERPOWERD_SESSION_TOKEN  set automatically by superpowers run-session
    SUPERPOWERD_SOCKET         override superpowerd socket path
    SUPERPOWERD_SOCKET_DIR     directory containing superpowerd.sock
"""
import sys
sys.path = sys.path[1:]  # strip script dir — may be a non-navigable fs (e.g. 9P)
import array
import json
import os
import signal
import socket
import struct
import threading

_token_lock = threading.Lock()
_LOG_PREFIX = "elevate-adapter-superpowerd"


class _SpPool:
    """Pool of pre-connected SEQPACKET sockets to superpowerd.

    Connections are created on demand and returned after use.
    Failed connections are discarded, not returned.
    """

    def __init__(self, max_idle=4):
        self._lock = threading.Lock()
        self._idle = []  # list of ready-to-use sockets
        self._max_idle = max_idle

    def acquire(self):
        """Get a connected socket, from pool or freshly created."""
        with self._lock:
            while self._idle:
                sp = self._idle.pop()
                # Verify connection is still alive (peek for hangup)
                try:
                    sp.setblocking(False)
                    try:
                        # If recv returns 0 or raises, connection is dead
                        data = sp.recv(1, socket.MSG_PEEK | socket.MSG_DONTWAIT)
                        if not data:
                            sp.close()
                            continue
                    except BlockingIOError:
                        # No data available = good, connection alive
                        pass
                    except OSError:
                        sp.close()
                        continue
                    finally:
                        sp.setblocking(True)
                    return sp
                except OSError:
                    continue

        # Pool empty — create a new connection
        sp = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET)
        sp.connect(sp_socket_path())
        return sp

    def release(self, sp):
        """Return a healthy connection to the pool."""
        with self._lock:
            if len(self._idle) < self._max_idle:
                self._idle.append(sp)
            else:
                sp.close()

    def close_all(self):
        """Close all pooled connections."""
        with self._lock:
            for sp in self._idle:
                sp.close()
            self._idle.clear()


_sp_pool = _SpPool()


def _log(msg):
    print(f"{_LOG_PREFIX}: {msg}", file=sys.stderr, flush=True)


def _read_master_token():
    """Read the superpowerd master token from the standard location."""
    path = os.environ.get("SUPERPOWERD_MASTER_TOKEN_FILE") or \
        os.path.join(os.path.expanduser("~"), "Private", "superpowerd_master.token")
    try:
        return open(path).read().strip()
    except OSError as e:
        _log(f"cannot read master token at {path}: {e}")
        return None


def refresh_token():
    """Create a fresh session directly on superpowerd via SEQPACKET.

    Connects to superpowerd and calls CreateSession with the master token.
    Because WE are the connecting process, the session PID = our PID,
    so pidAlive() stays true for our entire lifetime.

    Returns the new token string, or None on failure.
    Updates os.environ so subsequent calls use the fresh token.
    """
    _log("refreshing session token via direct CreateSession")
    master_token = _read_master_token()
    if not master_token:
        _log("refresh failed: cannot read master token")
        return None

    sp_path = sp_socket_path()
    try:
        sp = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET)
        sp.connect(sp_path)
    except OSError as e:
        _log(f"refresh failed: connect {sp_path}: {e}")
        return None

    try:
        req = json.dumps({
            "jsonrpc": "2.0",
            "method": "CreateSession",
            "id": 1,
            "params": {
                "auth": {"master_token": master_token},
                "env": dict(os.environ),
                "working_dir": os.getcwd(),
            },
        }).encode()
        sp.send(req)
        resp_data = sp.recv(1024 * 1024)
        sp.close()

        resp = json.loads(resp_data)
        if "error" in resp and resp["error"]:
            msg = resp["error"].get("message", str(resp["error"]))
            _log(f"refresh failed: superpowerd error: {msg}")
            return None

        token = resp.get("result", {}).get("token", "")
        if token:
            os.environ["SUPERPOWERD_SESSION_TOKEN"] = token
            _log("session token refreshed successfully")
            return token
        _log(f"refresh failed: no token in response: {resp_data!r:.200}")
    except (OSError, json.JSONDecodeError, ValueError) as e:
        _log(f"refresh failed: {e}")
        try:
            sp.close()
        except OSError:
            pass
    return None


def default_socket_path():
    xdg = os.environ.get("XDG_RUNTIME_DIR") or f"/run/user/{os.getuid()}"
    return os.path.join(xdg, "ollie", "elevate.sock")


def sp_socket_path():
    if p := os.environ.get("SUPERPOWERD_SOCKET"):
        return p
    if d := os.environ.get("SUPERPOWERD_SOCKET_DIR"):
        return d + "/superpowerd.sock"
    if x := os.environ.get("XDG_RUNTIME_DIR"):
        return x + "/superpowerd/superpowerd.sock"
    return f"/run/user/{os.getuid()}/superpowerd/superpowerd.sock"


def send_frame(conn, frame_type, payload):
    """Send a framed message: <type:1><length:4 big-endian><payload>"""
    data = payload if isinstance(payload, bytes) else payload.encode()
    header = frame_type.encode() + struct.pack(">I", len(data))
    conn.sendall(header + data)


def send_error(conn, msg, exit_code=1):
    """Send an error message and exit frame."""
    send_frame(conn, "d", msg + "\n")
    send_frame(conn, "x", str(exit_code))


def _is_auth_error(sp_data):
    """Check if superpowerd response is an authentication/session error."""
    if not sp_data:
        return False
    try:
        resp = json.loads(sp_data)
        err = resp.get("error")
        if isinstance(err, dict):
            msg = err.get("message", "").lower()
            return "invalid session token" in msg or "session not found" in msg or "session expired" in msg or "parent session expired" in msg
    except (json.JSONDecodeError, ValueError):
        pass
    return False


def _is_suspicious_failure(sp_data, output_received):
    """Check if the result looks like a silent session failure worth retrying.

    Exit 255 with no output is a strong signal that superpowerd rejected
    the command without running it (e.g. expired session, killed process).
    """
    if output_received:
        return False
    if not sp_data:
        return True
    try:
        resp = json.loads(sp_data)
        result = resp.get("result", {})
        exit_code = result.get("exit_code")
        if exit_code == 255:
            return True
    except (json.JSONDecodeError, ValueError):
        return True
    return False


def _execute_on_superpowerd(conn, cmd, cwd, token):
    """Execute a command via superpowerd. Returns (sp_data, True) on success/normal
    error, or (sp_data, False) if the caller should retry (auth error, no output streamed)."""
    # Acquire a connection from the pool
    try:
        sp = _sp_pool.acquire()
    except OSError as e:
        send_error(conn, f"elevate-adapter: connect {sp_socket_path()}: {e}")
        return None, True

    # Create pipes for stdout+stderr (merged) and stdin (/dev/null)
    stdin_r = os.open("/dev/null", os.O_RDONLY)
    stdout_r, stdout_w = os.pipe()

    # Build JSON-RPC request
    sp_req = json.dumps({
        "jsonrpc": "2.0",
        "method": "ExecuteCommandStreaming",
        "id": 1,
        "params": {
            "auth": {"session_token": token},
            "command": cmd,
            "working_dir": cwd,
            "approval_timeout": 7200,
        },
    }).encode()

    # Send request + 3 fds (stdin, stdout, stderr) via SCM_RIGHTS
    fds = array.array("i", [stdin_r, stdout_w, stdout_w])
    ancdata = [(socket.SOL_SOCKET, socket.SCM_RIGHTS, fds)]
    try:
        sp.sendmsg([sp_req], ancdata)
    except OSError as e:
        os.close(stdin_r)
        os.close(stdout_r)
        os.close(stdout_w)
        sp.close()
        send_error(conn, f"elevate-adapter: sendmsg: {e}")
        return None, True

    # Close write-ends in this process (superpowerd has them now)
    os.close(stdin_r)
    os.close(stdout_w)

    # Stream output from the pipe to the elevate client
    stream_done = threading.Event()
    client_broken = threading.Event()
    output_received = threading.Event()

    def stream_output():
        try:
            while True:
                chunk = os.read(stdout_r, 4096)
                if not chunk:
                    break
                output_received.set()
                send_frame(conn, "d", chunk)
        except (BrokenPipeError, OSError):
            client_broken.set()
        finally:
            os.close(stdout_r)
            stream_done.set()

    t = threading.Thread(target=stream_output, daemon=True)
    t.start()

    # Wait for superpowerd's response (comes after command exits)
    sp_ok = True
    try:
        sp_data = sp.recv(1024 * 1024)
    except OSError:
        sp_data = b""
        sp_ok = False

    # Return connection to pool if healthy, otherwise discard
    if sp_ok and sp_data:
        _sp_pool.release(sp)
    else:
        sp.close()

    # Wait for streaming thread to finish
    stream_done.wait()

    if client_broken.is_set():
        return None, True

    # If auth error or suspicious silent failure, signal retry
    if not output_received.is_set():
        if _is_auth_error(sp_data):
            _log(f"auth error detected, signalling retry (cmd={cmd!r:.60})")
            return sp_data, False
        if _is_suspicious_failure(sp_data, False):
            _log(f"suspicious failure (exit 255, no output), signalling retry (cmd={cmd!r:.60})")
            return sp_data, False

    return sp_data, True


def handle(conn):
    with conn:
        try:
            _handle_inner(conn)
        except Exception as e:
            try:
                send_error(conn, f"elevate-adapter: internal error: {e}")
            except OSError:
                pass


def _handle_inner(conn):
    data = b""
    while chunk := conn.recv(65536):
        data += chunk
        if b"\n" in data:
            break

    try:
        req = json.loads(data)
    except json.JSONDecodeError as e:
        send_error(conn, f"elevate-adapter: bad request: {e}")
        return

    cmd = req.get("cmd", "")
    cwd = req.get("cwd", os.getcwd())

    token = os.environ.get("SUPERPOWERD_SESSION_TOKEN", "")

    # First attempt
    sp_data, done = _execute_on_superpowerd(conn, cmd, cwd, token)
    if done:
        _finish_response(conn, sp_data)
        return

    # Auth/session error — refresh token and retry once
    _log(f"retrying after session failure (cmd={cmd!r:.60})")
    with _token_lock:
        # Re-read in case another thread already refreshed
        current = os.environ.get("SUPERPOWERD_SESSION_TOKEN", "")
        if current == token:
            new_token = refresh_token()
        else:
            _log("another thread already refreshed token, using new one")
            new_token = current

    if not new_token:
        _log(f"retry abandoned: token refresh failed (cmd={cmd!r:.60})")
        # Refresh failed — report the original auth error
        _finish_response(conn, sp_data)
        return

    sp_data, _ = _execute_on_superpowerd(conn, cmd, cwd, new_token)
    _finish_response(conn, sp_data)


def _finish_response(conn, sp_data):
    """Parse superpowerd response and send exit code frame to client."""
    if sp_data is None:
        return

    exit_code = 1
    try:
        sp_resp = json.loads(sp_data)
        rpc_err = sp_resp.get("error")
        if rpc_err:
            msg = rpc_err.get("message", str(rpc_err)) if isinstance(rpc_err, dict) else str(rpc_err)
            _log(f"superpowerd RPC error: {msg}")
            send_frame(conn, "d", f"elevate-adapter: {msg}\n".encode())
        else:
            result = sp_resp.get("result", {})
            exit_code = result.get("exit_code", 0)
    except (json.JSONDecodeError, ValueError):
        _log(f"superpowerd returned non-JSON response: {sp_data!r:.200}")

    send_frame(conn, "x", str(exit_code))


def main():
    if "SUPERPOWERD_SESSION_TOKEN" not in os.environ:
        # Get a token directly from superpowerd instead of os.execvp through
        # superpowers run-session (which fails on Gentoo's python-exec2c).
        token = refresh_token()
        if not token:
            _log("fatal: cannot obtain session token")
            sys.exit(1)

    sock_path = os.environ.get("OLLIE_ELEVATE_SOCKET") or default_socket_path()
    sock_path = os.path.expanduser(sock_path)
    os.makedirs(os.path.dirname(sock_path), exist_ok=True)

    if os.path.exists(sock_path):
        os.unlink(sock_path)

    srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    srv.bind(sock_path)
    srv.listen()

    def shutdown(sig, frame):
        srv.close()
        if os.path.exists(sock_path):
            os.unlink(sock_path)
        sys.exit(0)

    signal.signal(signal.SIGTERM, shutdown)
    signal.signal(signal.SIGINT, shutdown)

    _log(f"listening on {sock_path}")

    while True:
        try:
            conn, _ = srv.accept()
        except OSError:
            break
        threading.Thread(target=handle, args=(conn,), daemon=True).start()


if __name__ == "__main__":
    main()
