#!/usr/bin/env python3
"""Elevation plugin: run a command outside the sandbox via superpowers.

Usage: elevate -- <cmd>

Replace this script to use a different elevation mechanism.
"""
import json
import os
import socket
import sys


def _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 main():
    if "SUPERPOWERD_SESSION_TOKEN" not in os.environ:
        os.execvp("superpowers", ["superpowers", "run-session", "--",
                                   sys.executable, __file__] + sys.argv[1:])

    # sys.argv: ['elevate', '--', '<cmd>']
    if len(sys.argv) < 3:
        print("usage: elevate -- <cmd>", file=sys.stderr)
        sys.exit(1)

    cmd = sys.argv[2]
    cwd = os.getcwd()
    token = os.environ["SUPERPOWERD_SESSION_TOKEN"]

    sock_path = _socket_path()
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET)
    try:
        sock.connect(sock_path)
    except OSError as e:
        print(f"elevate: connect {sock_path}: {e}", file=sys.stderr)
        sys.exit(1)

    req = json.dumps({
        "jsonrpc": "2.0",
        "method": "ExecuteCommand",
        "id": 1,
        "params": {
            "auth": {"session_token": token},
            "command": cmd,
            "working_dir": cwd,
            "approval_timeout": 7200,
        },
    }).encode()

    with sock:
        sock.sendall(req)
        data = sock.recv(1024 * 1024)

    resp = json.loads(data)
    rpc_err = resp.get("error")
    if rpc_err:
        msg = rpc_err.get("message", str(rpc_err)) if isinstance(rpc_err, dict) else str(rpc_err)
        print(f"elevate: {msg}", file=sys.stderr)
        sys.exit(1)

    result = resp.get("result", {})
    out = result.get("stdout", "")
    err = result.get("stderr", "")
    if out:
        sys.stdout.write(out)
    if err:
        sys.stderr.write(err)
    sys.exit(result.get("exit_code", 0))


if __name__ == "__main__":
    main()
