#!/usr/bin/env python3
import asyncio
import atexit
import base64
import contextlib
import errno
import fcntl
import json
import os
import platform
import pty
import pwd
import random
import re
import shutil
import signal
import stat
import subprocess
import sys
import tempfile
import termios
import textwrap
import threading
import time
import traceback
import zipfile
from concurrent.futures import ThreadPoolExecutor

# pid -> Process
PROCESSES = {}
# List of pids that are completed. Their tasks can be awaited and removed.
COMPLETED = []
VERBOSE=0
LOGFILE=None
RUNLOOP=None
TASKS=[]
QUITTING=False
REGISTERED=[]
LASTPS={}
AUTOPOLL = 0
AUTOPOLL_TASK = None
TTY_TASK = None
RECOVERY_STATE={}
# 0: Not blocking on stdin
# 1: Blocking on stdin
# 2: Not blocking on stdin but send %ping before reading next command.
READSTATE=0
#{SUB}
BASEID=str(random.randint(0, 1048576)) + str(os.getpid()) + str(int(time.time() * 1000000))
IDCOUNT=0
SEARCH_TASK = None
# it2-over-ssh: the unix-socket server the remote `it2` connects to, and the
# accepted connections (connid -> asyncio StreamWriter). Dormant until iTerm2
# sends the "it2listen" command.
IT2_SERVER = None
IT2_CONNS = {}
# connids whose close iTerm2 initiated via it2close. The reader coroutine checks this in its
# finally so it does NOT echo a duplicate `%it2 close` back for a connid iTerm2 already tore
# down. A framer/client-initiated teardown (client disconnect, buffer-cap abort) is NOT in
# here, so it still emits close to inform iTerm2. Each entry is discarded by the reader.
IT2_LOCALLY_CLOSED = set()
# Path of the bound it2 socket, remembered so we can unlink it on shutdown (it
# lives in a persistent, per-user directory, so a leaked socket would linger).
IT2_SOCKET_PATH = None
# Liveness lock for the bound socket: a live framer holds an exclusive advisory flock on
# a sidecar "<socket>.lock" file for its entire lifetime. The kernel releases that lock
# when the process dies (including SIGKILL when an ssh connection drops), so another
# framer can tell a leaked socket from a live one by whether its lock is acquirable,
# with no timing race and without connecting to the socket. IT2_LOCK_FD is kept open the
# whole time precisely to hold the lock; IT2_LOCK_PATH is its file, unlinked on teardown.
IT2_LOCK_FD = None
IT2_LOCK_PATH = None
# Upper bound on a single it2 client's unsent (buffered) bytes. it2send writes without
# awaiting drain() -- draining on the shared command loop would let a stalled consumer's
# backpressure freeze send/kill/file/poll for the whole SSH session -- so instead we bound
# memory: a stalled consumer whose buffer grows past this is dropped, not the loop.
IT2_WRITE_BUFFER_MAX = 4 * 1024 * 1024  # 4 MiB
# framer emits one OSC 134 "%it2 <connid> data <b64>" frame per socket read (via
# send_esc), and iTerm2's conductor OSC parser accumulates a frame's payload one
# char at a time (O(n^2) in its length). Normal up-traffic is only tiny
# HELLO/CANCEL, so this cap is inert in practice; it just stops a misbehaving or
# hostile client on the socket from feeding one giant frame to that parser. The
# demux reassembles the it2.py wire frames across chunks.
IT2_READ_CHUNK = 1024

def squash(i):
    a = list(map(chr, list(range(48,58))+list(range(65,91))+list(range(97,123))))
    b = len(a)
    return a[i] if i < b else squash(i // b) + a[i % b]

def makeid():
    global IDCOUNT
    result = BASEID + str(IDCOUNT)
    IDCOUNT += 1
    return squash(int(result))

def log(message):
    pass

def send(q, data):
    if QUITTING:
        log("[squelched] " + str(data))
        return
    log("> " + str(data))
    q.append(data)

def send_esc(q, text):
    if QUITTING:
        log("[squelched] " + str(text))
        return
    log("> [osc 134] " + str(text) + " [st]")
    q.append('\033]134;:' + str(text) + '\033\\')

def lock():
    return []

def _write_all(fd, data):
    # os.write can write FEWER bytes than requested (a signal such as SIGWINCH delivered after
    # a partial write, or a full pipe). PEP 475 only auto-retries on EINTR when zero bytes were
    # written, so a partial write is NOT retried -- loop until every byte is out. Dropping the
    # tail here silently truncates an up-channel OSC frame and desyncs iTerm2's conductor
    # parser for the whole ssh session, not just one it2 command.
    if isinstance(data, str):
        data = data.encode('utf-8')
    view = memoryview(data)
    total = 0
    length = len(view)
    while total < length:
        try:
            written = os.write(fd, view[total:])
        except InterruptedError:
            continue  # EINTR with no progress: retry
        if written <= 0:
            break  # fd closed / unwritable; give up rather than spin
        total += written


def unlock(writes):
    if writes:
        for data in writes:
            _write_all(sys.stdout.fileno(), data)
        sys.stdout.flush()

class Process:
    @staticmethod
    async def run_tty(executable, args, cwd, env):
        master, slave = pty.openpty()
        try:
            def set_ctty(ctty_fd, master_fd):
                os.setsid()
                os.close(master_fd)
                fcntl.ioctl(ctty_fd, termios.TIOCSCTTY, 0)
                window_size = fcntl.ioctl(sys.stdin.fileno(), termios.TIOCGWINSZ, '00000000')
                fcntl.ioctl(ctty_fd, termios.TIOCSWINSZ, window_size)
            log(env)
            log(f'create subprocess args={args} cwd={cwd} executable={executable}')
            proc = await asyncio.create_subprocess_exec(
                *args,
                stdin=slave,
                stdout=slave,
                stderr=slave,
                cwd=cwd,
                env=env,
                executable=executable,
                preexec_fn=lambda: set_ctty(slave, master))
        finally:
            os.close(slave)
        return await Process.run_tty_proc(proc, master, f'run_tty({args})')

    @staticmethod
    async def run_shell_tty(command):
        master, slave = pty.openpty()
        try:
            env = dict(os.environ)
            env["LANG"] = "C"
            proc = await asyncio.create_subprocess_shell(
                command,
                stdin=slave,
                stdout=slave,
                stderr=slave,
                env=env)
        finally:
            os.close(slave)
        return await Process.run_tty_proc(proc, master, f'run_shell_tty({command})')

    @staticmethod
    async def run_tty_proc(proc, master, descr):
        pipe = open(master, 'wb', 0)
        writer = await Process._writer(pipe)
        reader, read_transport = await Process._reader(open(master, 'rb', 0))
        p = Process(proc, writer, reader, None, master=master, descr=descr)
        p.add_cleanup(read_transport.close)
        return p

    @staticmethod
    async def writer(fd):
        pipe = open(fd, 'wb', 0)

        return await Process._writer(pipe)

    @staticmethod
    async def _writer(pipe):
        loop = asyncio.get_event_loop()
        transport, protocol = await loop.connect_write_pipe(asyncio.Protocol, pipe)
        writer = asyncio.StreamWriter(
                transport=transport,
                protocol=protocol,
                reader=None,
                loop=loop)
        return writer

    @staticmethod
    async def reader(fd):
        pipe = open(fd, 'rb', 0)
        return Process._reader(pipe)

    @staticmethod
    async def _reader(pipe):
        loop = asyncio.get_event_loop()
        reader = asyncio.StreamReader()
        protocol = asyncio.StreamReaderProtocol(reader)
        transport, _ = await loop.connect_read_pipe(lambda: protocol, pipe)
        return reader, transport

    def __init__(self, process, writer, stdout_reader, stderr_reader, master=None, descr=""):
        self.__process = process
        self.__writer = writer
        self.__stdout_reader = stdout_reader
        self.__stderr_reader = stderr_reader
        self.__stdout_read_handler = None
        self.__stderr_read_handler = None
        self.__cleanup = []
        self.__return_code = None
        self.__master = master
        self.__descr = descr
        self.echo = True
        self.icanon = True
        self.login = False

    @property
    def master(self):
        return self.__master

    async def cleanup(self):
        log(f'cleanup {self.__descr}: cleanup process {self.pid}')
        if self.__return_code is None:
            self.__return_code = self.__process.returncode
            log(f'nonblocking return code is {self.__return_code}')
        if self.__return_code is None:
            log(f'kill {self.__descr}')
            try:
                # Use os.kill directly to avoid asyncio bug with send_signal on dead processes
                # https://github.com/python/cpython/issues/87744
                # https://bugs.python.org/issue38630
                os.kill(self.pid, signal.SIGKILL)
            except ProcessLookupError:
                log(f'cleanup {self.__descr}: process already dead')
            except Exception as e:
                log(f'cleanup {self.__descr}: exception {e} during kill')
            log(f'cleanup {self.__descr}: wait')
            await self.wait()
        log(f'cleanup {self.__descr}: close writer tx')
        self.__writer.transport.close()
        log(f'cleanup {self.__descr}: adding handlers to TASKS')
        global TASKS
        if self.__stderr_read_handler:
            TASKS.append(self.__stderr_read_handler)
        TASKS.append(self.__stdout_read_handler)
        log(f'cleanup {self.__descr}: running cleanup callbacks')

        # StreamReader doesn't have a transport so it must be closed by a __cleanup function.
        for f in self.__cleanup:
            f()
        log(f'cleanup {self.__descr}: done')

    def add_cleanup(self, coro):
        self.__cleanup.append(coro)

    async def kill(self, signal):
        self.__process.send_signal(signal)

    async def wait(self):
        rc = await self.__process.wait()
        self.__return_code = rc
        return rc

    @property
    def return_code(self):
        return self.__return_code

    @property
    def pid(self):
        return self.__process.pid

    async def readline(self):
        return await self.__stdout_reader.readline()

    async def read_forever(self, reader, channel, callback):
        try:
            while True:
                log(f'read_forever {self.__descr}: reading for channel {channel}')
                value = await reader.read(256)
                log(f'read_forever {self.__descr}: read {value} for channel {channel}')
                coro = callback(channel, value)
                if coro:
                    log(f'read_forever {self.__descr}: await callback-returned coro {coro}')
                    await coro
                if len(value) == 0:
                    return
        except IOError:
            log(f'read_forever {self.__descr}: stopping because of IOError')
            coro = callback(channel, b'')
            await coro
            return
        except Exception as e:
            log(f'read_forever {self.__descr}: {e}')

    async def handle_read(self, callback):
        self.__stdout_read_handler = asyncio.create_task(self.read_forever(self.__stdout_reader, 1, callback))
        if self.__stderr_reader:
            self.__stderr_read_handler = asyncio.create_task(self.read_forever(self.__stderr_reader, 2, callback))

    async def write(self, data):
        self.__writer.write(data)

    def send_signal(self, signal):
        self.__process.send_signal(signal)

## Login Shell

def guess_login_shell():
    path = pwd.getpwuid(os.geteuid()).pw_shell
    if os.access(path, os.X_OK):
        return path
    return "/bin/sh"

## Process monitoring

async def autopoll(delay):
    try:
        global AUTOPOLL
        while True:
            log('autopoll: call poll()')
            cats = await poll()
            if not cats:
                log(f'autopoll: sleep for {delay}')
                await asyncio.sleep(delay)
                log(f'autopoll: awoke')
                continue
            # Send poll output and sleep until client requests autopolling again.
            identifier = makeid()
            q = lock()
            send_esc(q, f'%autopoll {identifier}')
            send_poll_output(q, cats)
            send_esc(q, f'%end {identifier}')
            unlock(q)
            AUTOPOLL = 0
            while not AUTOPOLL:
                log(f'autopoll: sleep for {delay}')
                await asyncio.sleep(delay)
                log(f'autopoll: awoke')
    except asyncio.CancelledError:
        log('autopoll canceled')
        raise
    except Exception as e:
        log(f'autopoll threw {e}: {traceback.format_exc()}')

def send_poll_output(q, cats):
    for catname in cats:
        output = cats[catname]
        send(q, f'$begin {catname}'.encode("utf-8") + b'\n')
        for line in output:
            send(q, line.encode("utf-8") + b'\n')

def get_echo_icanon(tty):
    try:
        attrs = termios.tcgetattr(tty)
        return (bool(attrs[3] & termios.ECHO), bool(attrs[3] & termios.ICANON))
    except Exception as e:
        log(f'get_echo_icanon threw {e}: {traceback.format_exc()}')
        raise

async def watch_tty(proc, delay):
    try:
        while True:
            log("watch_tty: poll")
            poll_tty(proc)
            await asyncio.sleep(delay)
    except asyncio.CancelledError:
        log('watch_tty canceled')
        raise
    except Exception as e:
        log(f'watch_tty threw {e}: {traceback.format_exc()}')

def poll_tty(proc):
    log(f"Check TTY with fd {proc.master}")
    new_echo, new_icanon = get_echo_icanon(proc.master)
    if new_echo != proc.echo or new_icanon != proc.icanon:
        log(f'echo: {proc.echo}->{new_echo}, icanon: {proc.icanon}->{new_icanon}')
        parts = [
            ('+' if new_echo else '-') + 'echo',
            ('+' if new_icanon else '-') + 'icanon']
        print_tty(" ".join(parts))
    proc.echo = new_echo
    proc.icanon = new_icanon

async def poll():
    result = {}
    ps_out = await poll_ps()
    if ps_out is not None:
        result["ps"] = ps_out
    cpu_time_diff = await poll_cpu()
    if cpu_time_diff is not None:
        result["cpu"] = cpu_time_diff
    return result

mpstat_exists = None  # Global variable to cache the mpstat existence check

async def check_mpstat_exists():
    global mpstat_exists
    if mpstat_exists is None:
        try:
            log("Checking if mpstat exists")
            proc = await asyncio.create_subprocess_shell(
                "mpstat -V",
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            output, _ = await proc.communicate()
            log(f'return code is {proc.returncode}')
            mpstat_exists = proc.returncode == 0
        except FileNotFoundError:
            log("mpstat: file not found")
            mpstat_exists = False


async def poll_cpu():
    operating_system = platform.system()
    if operating_system == 'Darwin':  # macOS
        command = "top -l 1 -n 0 | awk '/CPU usage/ {print $3}'"
    elif operating_system == 'Linux':  # Linux
        await check_mpstat_exists()
        if not mpstat_exists:
            return None
        command = "mpstat -P ALL 1 1 | awk '/Average:/ && $2 == \"all\" {print 100 - $NF}'"
    else:
        return None
    env = dict(os.environ)
    env["LANG"] = "C"
    proc = await asyncio.create_subprocess_shell(
        command,
        stdin=asyncio.subprocess.PIPE,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        env=env)
    log(f'poll_cpu: run {command} {proc}')
    output, erroutput = await proc.communicate()
    if proc.returncode == 0:
        log(f'poll_cpu: successful return')
        # Parse the output here
        final = "=" + output.decode("utf-8").strip()
        log(f'poll_cpu: return parsed output')
        return [final]
    log(f'poll_cpu: {command} failed with {proc.returncode}')
    return None

async def poll_ps():
    env = dict(os.environ)
    env["LANG"] = "C"
    proc = await asyncio.create_subprocess_shell(
        "ps -eo pid,ppid,stat,lstart,command",
        stdin=asyncio.subprocess.PIPE,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        env=env)
    log(f'poll_ps: run ps {proc}')
    output, erroutput = await proc.communicate()
    if proc.returncode == 0:
        log(f'poll_ps: successful return')
        final = procmon_parse(output)
        log(f'poll_ps: return parsed output')
        return final
    log(f'poll_ps: ps failed with {proc.returncode}')
    return None

async def register(pid):
    global REGISTERED
    if pid in REGISTERED:
        return
    REGISTERED.append(pid)
    log(f'After registering {pid} REGISTERED={REGISTERED}')

async def deregister(pid):
    global REGISTERED
    if pid in REGISTERED:
        REGISTERED.remove(pid)

def procmon_parse(output):
    output = output.decode("utf-8")
    lines = output.split("\n")
    whitespace = r'\s+'
    number = r'\d+'
    nonspace = r'\S+'
    letters = r'[A-Za-z]+'
    pattern = "".join(
        [r'^',
         r'\s*',
         r'(',
         number,  # pid [capture 1]
         r')',
         whitespace,
         r'(',
         number,  # ppid [capture 2]
         r')',
         whitespace,
         r'(',
         nonspace,  # stat [capture 3]
         r')',
         whitespace,
         r'(',
         letters,  # day of week  [capture 4]
         whitespace,
         letters,  # name of month
         whitespace,
         number,  # day of month
         whitespace,
         number,  # hh
         r':',
         number,  # mm
         r':',
         number,  # ss
         whitespace,
         number,  # yyyy
         r')',
         whitespace,
         r'(.*)'  # command  [capture 5]
    ])
    def parse(line):
        match = re.search(pattern, line)
        if not match:
            return None
        return (match.group(1), match.group(2), match.group(3), match.group(4), match.group(5))
    rows = map(parse, lines)
    # pid->ppid
    parent={}
    # ppid->[pid]
    children={}
    # pid->row
    index={}
    for row in rows:
        if row is None:
            continue
        if row[4].startswith("(") and row[4].endswith(")"):
            log(f'procmon_parse: ignore defunct {row}')
            continue
        pid = row[0]
        ppid = row[1]
        parent[pid] = ppid
        children[ppid] = children.get(ppid, []) + [pid]
        index[pid] = row
    log(f'procmon_parse: {len(index)} valid rows')
    # pid -> row
    results = {}
    def add(pid):
        results[pid] = index[pid]
        for child in children.get(pid, []):
            add(child)
    for pid in REGISTERED:
        log(f'procmon_parse: Add hierarchy starting at {pid}')
        if str(pid) in index:
            add(str(pid))
    log(f'procmon_parse: {len(results)} processes in output')

    global LASTPS
    last = dict(LASTPS)
    LASTPS = dict(results)

    def diff():
        currentkeys = set(results.keys())
        lastkeys = set(last.keys())
        log(f'procmon_parse: diff current={currentkeys} last={lastkeys}')
        for addition in currentkeys - lastkeys:
            log(f'procmon_parse: add {addition}')
            yield "+" + " ".join(map(str, results[addition]))
        for removal in lastkeys - currentkeys:
            log(f'procmon_parse: remove {removal}')
            yield "-" + str(removal)
        for pid in results:
            if pid in last and results[pid] != last[pid]:
                log(f'procmon_parse: edit {pid}')
                yield "~" + " ".join(map(str, results[pid]))
    return list(diff())


async def get_env_var(var_name):
    # Get the user's shell
    user_shell = os.environ.get('SHELL', '/bin/bash')

    # Non-POSIX shells may leave POSIX shell path in the $SHELL; perform an additional check.
    if os.environ.get('XONSHRC', ''):
        user_shell = 'xonsh'

    known_shells = ['bash', 'csh', 'dash', 'fish', 'ksh', 'sh', 'tcsh', 'xonsh', 'zsh']

    # Extract the last path component of the user's shell
    shell_name = os.path.basename(user_shell)
    log(f'Shell is {user_shell}')

    if shell_name not in known_shells:
        log(f'Unknown shell {shell_name}')
        return None

    command = f'echo ${var_name}'

    try:
        process = await asyncio.create_subprocess_exec(
            user_shell, '-c', command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )

        stdout, stderr = await process.communicate()

        # Print the value of the environment variable
        return stdout.decode().strip()
    except Exception as e:
        log(f'get_env_var({var_name}) threw: {traceback.format_exc()}')

# Run the async function
async def save_and_exec(identifier, code):
    log("begin save_and_exec")
    try:
        tmp = tempfile.NamedTemporaryFile(delete=False)
        tmpname = tmp.name
        log(f'save_and_exec: tmp={tmpname}')
        tmp.write(code.encode("utf-8"))
        tmp.close()
        st = os.stat(tmpname)
        os.chmod(tmpname, st.st_mode | stat.S_IEXEC)

        e = dict(os.environ)
        e["SELF"] = tmpname
        await run_login_shell(identifier, "/bin/sh", ["sh", tmpname], os.getcwd(), e)
    except Exception as e:
        log(f'save_and_exec: {traceback.format_exc()}')
        q = begin(identifier)
        send_esc(q, f'Command failed: {e}')
        end(q, identifier, 1)

## Commands

async def handle_login(identifier, args):
    log("begin handle_login")
    cwd = args[0]
    args = args[1:]
    log(f'cwd={cwd} args={args}')
    cwd = os.path.expandvars(os.path.expanduser(cwd))
    had_args = len(args) > 0
    if args:
        args = ["-c", " ".join(args)]
        log(f'Update args to {args}')
    login_shell = guess_login_shell()
    _, shell_name = os.path.split(login_shell)
    argv0 = "-" + shell_name

    return await run_login_shell(identifier, login_shell, [argv0] + args, cwd, os.environ)

async def run_login_shell(identifier, login_shell, argv, directory, environment):
    log(f'run_login_shell: dir={directory} command={login_shell} argv={argv}')
    try:
        proc = await Process.run_tty(
            login_shell,
            argv,
            directory,
            environment)
        log("login shell started")
        global PROCESSES
        proc.login = True
        PROCESSES[proc.pid] = proc
        q = begin(identifier)
        send_esc(q, proc.pid)
        end(q, identifier, 0)
        await proc.handle_read(make_monitor_process(proc, True))
    except Exception as e:
        log(f'run_login_shell: {e}')
        q = begin(identifier)
        send_esc(q, f'Command failed: {e}')
        end(q, identifier, 1)
    return False

async def handle_run(identifier, args):
    """Run a command inside the user's shell"""
    try:
        proc = await Process.run_shell_tty(args[0])
        global PROCESSES
        PROCESSES[proc.pid] = proc
        q = begin(identifier)
        send_esc(q, proc.pid)
        end(q, identifier, 0)
        await proc.handle_read(make_monitor_process(proc, False))

        start_tty_task(identifier, proc)
    except Exception as e:
        log(f'handle_run: {e}')
        q = begin(identifier)
        end(q, identifier, 1)
    return False

def start_tty_task(identifier, proc):
    proc.tty_task = asyncio.create_task(watch_tty(proc, 1))

def reset():
    global REGISTERED
    global LASTPS
    global AUTOPOLL
    REGISTERED = []
    LASTPS = {}
    AUTOPOLL = 0

async def handle_reset1(identifier, args):
    reset()
    q = begin(identifier)
    end(q, identifier, 0)

async def handle_reset2(identifier, args):
    reset()
    q = begin(identifier)
    send_esc(q, args[0])
    end(q, identifier, 0)

async def handle_save(identifier, args):
    log(f'handle_save {identifier} {args}')
    global RECOVERY_STATE
    try:
        RECOVERY_STATE = dict(s.split('=', 1) for s in args)
        code = 0
    except Exception as e:
        log(f'handle_save: {e}')
        code = 1
    q = begin(identifier)
    end(q, identifier, code)

async def handle_getenv(identifier, args):
    log(f'handle_getenv {identifier} {args}')
    q = begin(identifier)
    value = await get_env_var(args[0])
    if q:
        send_esc(q, value)
    end(q, identifier, 0)

async def handle_runpy(identifier, args):
    "args[0] should be a sequence of Python statements returning a string"
    log(f'handle_runpy {identifier} {args}')
    q = begin(identifier)
    try:
        encoded = args[0]
        script = base64.b64decode(encoded).decode('utf-8')
        indented = textwrap.indent(script, '    ')
        wrapper = f"def __f__():\n{indented}"
        env = {}
        exec(wrapper, env)
        output = env['__f__']()
        send_esc(q, str(output) if output is not None else "")
        end(q, identifier, 0)
    except:
        send_esc(q, traceback.format_exc())
        end(q, identifier, 1)

async def handle_eval(identifier, args):
    log(f'handle_eval {identifier} [{len(args[0])} bytes]')
    await save_and_exec(identifier, base64.b64decode(args[0]).decode('latin1'))

def make_path(b64path):
    return os.path.expanduser(base64.b64decode(b64path)).decode('latin1')

async def handle_file(identifier, args):
    log(f'handle_file {identifier} {args}')
    if len(args) < 2:
        log(f'Not enough args {args}')
        q = begin(identifier)
        end(q, identifier, 1)
        return
    q = begin(identifier)
    sub = args[0]
    log(f'sub is {sub}')
    if sub == "ls":
        await handle_file_ls(q, identifier, make_path(args[1]), args[2])
        return
    if sub == "fetch":
        if len(args) >= 4:
            await handle_file_fetch(q, identifier, make_path(args[1]), int(args[2]), int(args[3]))
        else:
            await handle_file_fetch(q, identifier, make_path(args[1]), 0, float('inf'))
        return
    if sub == "stat":
        await handle_file_stat(q, identifier, make_path(args[1]))
        return
    if sub == "suggest":
        await handle_file_suggest(q,
                                  identifier,
                                  make_path(args[1]),
                                  args[2],
                                  make_path(args[3]),
                                  args[4],
                                  int(args[5]))
        return
    if sub == "rm":
        i = 1
        recursive = False
        while i < len(args) and args[i].startswith("-"):
            if args[i] == "-r":
                recursive = True
            i += 1
        if i == len(args):
            end(q, identifier, 1)
            return
        await handle_file_rm(q, identifier, make_path(args[i]), recursive)
        return
    if sub == "ln":
        await handle_file_ln(
            q,
            identifier,
            make_path(args[1]),
            make_path(args[2]))
        return
    if sub == "mv":
        await handle_file_mv(
            q,
            identifier,
            make_path(args[1]),
            make_path(args[2]))
        return
    if sub == "mkdir":
        await handle_file_mkdir(q, identifier, make_path(args[1]))
        return
    if sub == "create":
        await handle_file_create(q,
                                 identifier,
                                 make_path(args[1]),
                                 base64.b64decode("".join(args[2:])))
        return
    if sub == "append":
        await handle_file_append(q,
                                 identifier,
                                 make_path(args[1]),
                                 base64.b64decode("".join(args[2:])))
        return
    if sub == "utime":
        await handle_file_utime(q,
                                identifier,
                                make_path(args[1]),
                                int(float(args[2])))
        return
    if sub == "chmod-u":
        await handle_file_chmod_u(q,
                                  identifier,
                                  make_path(args[1]),
                                  args[2])
        return
    if sub == "zip":
        await handle_file_zip(q,
                              identifier,
                              make_path(args[1]))
        return
    if sub == "search":
        await handle_file_search(q,
                                 identifier,
                                 args[1:])
        return

    log(f'unrecognized subcommand {sub}')
    end(q, identifier, 1)

def permissions(path):
    return {
        "r": os.access(path, os.R_OK, effective_ids=True),
        "w": os.access(path, os.W_OK, effective_ids=True),
        "x": os.access(path, os.X_OK, effective_ids=True) }

def remotefile(pp, abspath, s):
    if stat.S_ISLNK(s.st_mode):
        k = {"symlink": {"_0": os.readlink(abspath) }}
    elif stat.S_ISDIR(s.st_mode):
        k = {"folder": {}}
    elif stat.S_ISREG(s.st_mode):
        k = {"file": {"_0": {"size": s[stat.ST_SIZE]}}}
    else:
        return None
    to = 978307200
    ctime = s[stat.ST_CTIME] - to
    mtime = s[stat.ST_MTIME] - to
    return {"absolutePath": abspath,
            "kind": k,
            "permissions": permissions(abspath),
            "parentPermissions": pp,
            "ctime": ctime,
            "mtime": mtime}

def file_error(q, identifier, e, path):
    if e is None:
        log(f'file_error {path}')
        end(q, identifier, 1)
        return
    try:
        raise e
    except OSError as e:
        log(f'file_error {path}: {traceback.format_exc()}')
        errors = {errno.EPERM: 1, errno.ENOENT: 2, errno.ENOTDIR: 3, errno.ELOOP: 4}
        end(q, identifier, errors.get(e.errno, 100))
    except PermissionError as e:
        log(f'file_error {path}: {traceback.format_exc()}')
        end(q, identifier, 1)
    except Exception as e:
        log(f'file_error {path}: {traceback.format_exc()}')
        end(q, identifier, 255)

async def handle_file_ls(q, identifier, path, sorting):
    log(f'handle_file_ls {identifier} {path}')
    try:
        pp = permissions(path)
        files = [(os.path.join(path, f),
                  os.stat(os.path.join(path, f), follow_symlinks=False))
                  for f in os.listdir(path)]
        def fmt(t):
            return remotefile(pp, t[0], t[1])

        obj = list(filter(lambda x: x is not None, map(fmt, files)))
        def sort_ls(d):
            if sorting == 'n':
                return d['absolutePath']
            return d['mtime']
        obj = sorted(obj, key=sort_ls)
        log(f'After sorting contents of {path} are {obj}')
        send_esc(q, "[")
        first = True
        for entry in obj:
            if first:
                first = False
            else:
                send_esc(q, ",")
            j = json.dumps(entry)
            send_esc(q, j)
        send_esc(q, "]")
        end(q, identifier, 0)
        log("handle_file_ls completed normally")
    except Exception as e:
        file_error(q, identifier, e, path)

def remote_file_json(path):
    pp = permissions(os.path.abspath(os.path.join(path, os.pardir)))
    s = os.stat(path, follow_symlinks=False)
    obj = remotefile(pp, path, s)
    return json.dumps(obj)

def send_remote_file(q, path):
    send_esc(q, remote_file_json(path))

async def handle_file_stat(q, identifier, path):
    log(f'handle_file_stat {identifier} {path}')
    try:
        send_remote_file(q, path)
        end(q, identifier, 0)
    except Exception as e:
        file_error(q, identifier, e, path)

async def handle_file_suggest(q, identifier, prefix, directories, pwd, permissions, max_count):
    log(f'handle_file_suggest({identifier}, {prefix}, {directories}, {pwd}, {permissions})')
    combined = []
    for relative_directory in decode_base64_directories(directories):
        if relative_directory.startswith('/'):
            log(f'relative directory is absolute {relative_directory}')
            directory = relative_directory
        else:
            log(f'relative directory is relative {relative_directory}')
            if not pwd:
                log(f'no pwd {pwd}')
                continue
            directory = os.path.join(pwd, relative_directory)
            log(f'joined directory is {directory}')
        temp = await really_find_completions_with_prefix(prefix, directory, max_count, 'x' in permissions)
        temp.sort()
        combined.extend([entry.removeprefix(prefix) for entry in temp])
        if len(combined) > max_count:
            break
    log(f'Return suggestions: {combined}')
    send_esc(q, json.dumps(combined))
    end(q, identifier, 0)

def decode_base64_directories(encoded_directories):
    decoded_directories = []
    directories = encoded_directories.split()

    for directory in directories:
        decoded_directory = base64.b64decode(directory).decode('utf-8')
        decoded_directories.append(decoded_directory)

    return decoded_directories

async def really_find_completions_with_prefix(prefix, directory, max_count, executable):
    log(f'really_find_completions_with_prefix({prefix}, {directory}, {max_count}, {executable}')
    if not prefix.startswith('/') and directory.startswith('/'):
        log('Prefix is relative and directory is absolute. Recurse.')
        temp = await really_find_completions_with_prefix(
            os.path.join(directory, prefix), '', max_count, executable)
        prefix_to_remove = directory if directory.endswith('/') else directory + '/'
        return [entry.removeprefix(prefix_to_remove) for entry in temp]

    if prefix.endswith('/'):
        log('Prefix ends in slash')
        return await contents_of_directory(prefix, '', executable, max_count)

    results = []
    is_directory = False
    exists = os.path.exists(prefix)
    if exists:
        is_directory = os.path.isdir(prefix)
        if is_directory:
            results.append(prefix + '/')

    container = os.path.dirname(prefix)
    if len(container) == 0:
        log(f"No dirname for {prefix}")
        return results

    log("Add contents of directory using prefix basename")
    results.extend(await contents_of_directory(container, os.path.basename(prefix), executable, max_count))
    return results

async def contents_of_directory(directory, prefix, executable, max_count):
    log(f'contents_of_directory({directory}, {prefix}, {executable}, {max_count}')
    try:
        relative = await asyncio.get_event_loop().run_in_executor(None, lambda: os.listdir(directory))
        result = []
        for path in relative:
            if len(result) >= max_count:
                break
            file_name = os.path.basename(path)
            if len(prefix) == 0 or file_name.startswith(prefix):
                full_path = os.path.join(directory, path)
                if not executable or os.access(full_path, os.X_OK):
                    result.append(full_path)
        return result
    except Exception as e:
        log(f'exception while getting contents of {directory}: {traceback.format_exc()}')
        return []

async def handle_file_zip(q, identifier, path):
    log(f'handle_file_zip {identifier} {path}')

    if not os.path.isdir(path):
        log(f'{path} is not a directory')
        file_error(q, identifier, None, f'{path} is not a directory')
        return
    try:
        # Create a temporary file in the user's home directory
        temp_file = tempfile.NamedTemporaryFile(dir=os.path.expanduser("~"), prefix=".", delete=False)
        log(f'zip to {temp_file}')
        ziph = zipfile.ZipFile(temp_file.name, 'w', zipfile.ZIP_DEFLATED)
        log('Opened zip for writing')
    except Exception as e:
        file_error(q, identifier, e, "While creating a temporary file under your home directory")
        return

    try:
        for root, dirs, files in os.walk(path):
            for file in files:
                try:
                    full_path = os.path.join(root, file)
                    log('add {full_path}')
                    ziph.write(full_path,
                               os.path.relpath(os.path.join(root, file),
                                               os.path.join(path, '..')))
                except Exception as e:
                    log(f'Failed to add {file}: {e}')
                    # zip as much as possible. Errors are just ignored, which isn't great.
                    pass
        log('Done zipping')
    except Exception as e:
        file_error(q, identifier, e, temp_file)
    finally:
        log('close zip')
        ziph.close()
        log('send tempfile name')
        send_esc(q, temp_file.name)
        end(q, identifier, 0)

class Search:
    def __init__(self, query: str, basedir: str):
        self.query = query
        self.basedir = basedir
        self.id = makeid()
        self.canceled = False

        # allow up to `window_size` in-flight notifications
        self.window_size = 16
        self.sema = threading.BoundedSemaphore(self.window_size)

        self._executor = ThreadPoolExecutor(
            max_workers=1,
            thread_name_prefix=f"Search-{self.id}"
        )
        self.task = asyncio.create_task(self.mainloop())
        self.loop: asyncio.AbstractEventLoop

    def emit(self, json: str) -> None:
        q = lock()
        send_esc(q, f'%notif search {self.id} {json}')
        unlock(q)

    def ack(self, count: int) -> None:
        # release up to `count` permits (never exceed window_size)
        for _ in range(count):
            try:
                self.sema.release()
            except ValueError:
                break

    def cancel(self) -> None:
        self.canceled = True
        # ensure any blocked acquire unblocks
        try:
            self.sema.release()
        except ValueError:
            pass

    async def mainloop(self) -> None:
        self.loop = asyncio.get_running_loop()
        try:
            await self.loop.run_in_executor(self._executor, self._scan_sync)
        except Exception as e:
            log(f"[Search {self.id}] error: {e!r}")
        finally:
            self._executor.shutdown(wait=False)

    def _scan_sync(self) -> None:
        stack = [self.basedir]
        while stack and not self.canceled:
            path = stack.pop()
            try:
                log(f"Scan {path} for {self.query}")
                with os.scandir(path) as it:
                    for entry in it:
                        if self.canceled:
                            return
                        try:
                            if entry.is_dir(follow_symlinks=False):
                                stack.append(entry.path)
                            elif entry.is_file(follow_symlinks=False) and self.query in entry.name:
                                # block here if we've hit window_size
                                self.sema.acquire()
                                if self.canceled:
                                    return
                                log(f"Found match: {entry.name} for {self.query}")
                                json = remote_file_json(entry.path)
                                self.loop.call_soon_threadsafe(self.emit, json)
                        except OSError:
                            continue
            except OSError:
                continue

async def handle_file_search(q, identifier, args):
    """Operate on searches, which run in the background.

    search start <query> <base dir>
    search ack <id> <count>
    search stop <id>
    """
    log(f'handle_file_search {identifier} {args}')
    operation = args[0]
    global SEARCH_TASK
    if operation == "start":
        query = args[1]
        basedir = make_path(args[2])
        if SEARCH_TASK:
            SEARCH_TASK.cancel()
            SEARCH_TASK = None
        SEARCH_TASK = Search(base64.b64decode(query).decode('utf-8'), basedir)
        send_esc(q, SEARCH_TASK.id)
    elif operation == "ack":
        id = args[1]
        count = int(args[2])
        if SEARCH_TASK and SEARCH_TASK.id == id:
            SEARCH_TASK.ack(count)
    elif operation == "stop":
        id = args[1]
        if SEARCH_TASK and SEARCH_TASK.id == id:
            if SEARCH_TASK:
                SEARCH_TASK.cancel()
                SEARCH_TASK = None
    end(q, identifier, 0)

async def handle_file_fetch(q, identifier, path, offset, size):
    log(f'handle_file_fetch {identifier} {path}')
    try:
        with open(path, "rb") as f:
            if offset > 0:
                f.seek(offset)
            if size == float('inf'):
                content = f.read()
            else:
                content = f.read(size)
            log(type(content))
            encoded = base64.encodebytes(content).decode('utf8')
            for line in encoded.split("\n"):
                log(f'will send f{type(line)}')
                send_esc(q, line)
        log("ending")
        end(q, identifier, 0)
    except Exception as e:
        file_error(q, identifier, e, path)

async def handle_file_rm(q, identifier, path, recursive):
    log(f'handle_file_rm {identifier} {path} {recursive}')
    try:
        if os.path.isdir(path):
            if recursive:
                shutil.rmtree(path)
            else:
                os.rmdir(path)
        else:
            os.unlink(path)
        log("ending")
        end(q, identifier, 0)
    except Exception as e:
        file_error(q, identifier, e, path)

async def handle_file_ln(q, identifier, pointTo, symlink):
    log(f'handle_file_ln {pointTo} {symlink}')
    try:
        os.symlink(pointTo, symlink)
        send_remote_file(q, symlink)
        end(q, identifier, 0)
    except Exception as e:
        file_error(q, identifier, e, symlink)

async def handle_file_mv(q, identifier, source, dest):
    log(f'handle_file_mv {source} {dest}')
    try:
        shutil.move(source, dest)
        send_remote_file(q, dest)
        end(q, identifier, 0)
    except Exception as e:
        file_error(q, identifier, e, source)

async def handle_file_mkdir(q, identifier, path):
    log(f'handle_file_mkdir {identifier} {path}')
    try:
        os.mkdir(path)
        end(q, identifier, 0)
    except Exception as e:
        file_error(q, identifier, e, path)

async def handle_file_create(q, identifier, path, content):
    log(f'handle_file_create {identifier} {path} length={len(content)} bytes')
    try:
        with open(path, "wb") as f:
            f.write(content)
        send_remote_file(q, path)
        end(q, identifier, 0)
    except Exception as e:
        file_error(identifier, e, path)

async def handle_file_append(q, identifier, path, content):
    log(f'handle_file_append {identifier} {path} length={len(content)} bytes')
    try:
        with open(path, "ab") as f:
            f.write(content)
        send_remote_file(q, path)
        end(q, identifier, 0)
    except Exception as e:
        file_error(q, identifier, e, path)

async def handle_file_utime(q, identifier, path, date):
    log(f'handle_file_utime {identifier} {path} {date}')
    try:
        os.utime(path, (date, date))
        send_remote_file(q, path)
        end(q, identifier, 0)
    except Exception as e:
        file_error(q, identifier, e, path)

async def handle_file_chmod_u(q, identifier, path, mode):
    log(f'handle_file_chmod_u {identifier} {path} {mode}')
    try:
        s = os.stat(path, follow_symlinks=False)
        value = s.st_mode & ~(stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR)
        if "r" in mode:
            value |= stat.S_IRUSR
        if "w" in mode:
            value |= stat.S_IWUSR
        if "x" in mode:
            value |= stat.S_IXUSR
        os.chmod(path, value)
        send_remote_file(q, path)
        end(q, identifier, 0)
    except Exception as e:
        file_error(q, identifier, e, path)

async def handle_recover(identifier, args):
    log("handle_recover")
    q = lock()
    # Keep this section await-free through unlock(q): the begin-recovery..end-recovery
    # block must stay contiguous. _it2_close_connections() aborts orphaned it2 connections
    # whose reader coroutines emit a `%it2 close` in their finally; abort() defers that via
    # call_soon so it runs only after we return, but a stray await here could let a woken
    # reader interleave that line into the recovery block.
    reset()
    _it2_close_connections()
    send_esc(q, f'begin-recovery')
    send_esc(q, f'recovery: version 2')
    for pid in PROCESSES:
        proc = PROCESSES[pid]
        if proc.login:
            send_esc(q, f'recovery: login {pid}')
            for key in RECOVERY_STATE:
                send_esc(q, f'recovery: {key} {RECOVERY_STATE[key]}')
        else:
            send_esc(q, f'recovery: process {pid}')
    send_esc(q, f'end-recovery')
    unlock(q)

async def handle_register(identifier, args):
    if len(args) < 1:
        fail(identifier, "not enough arguments")
        return
    try:
        pid = int(args[0])
    except Exception as e:
        log(f'handle_register: {e}')
        fail(identifier, "exception decoding argument")
        return
    q = begin(identifier)
    end(q,identifier, 0)
    await register(pid)

async def handle_deregister(identifier, args):
    if len(args) < 1:
        fail(identifier, "not enough arguments")
        return
    try:
        pid = int(args[0])
    except Exception as e:
        log(f'handle_deregister: {e}')
        fail(identifier, "exception decoding argument")
        return
    q = begin(identifier)
    end(q,identifier, 0)
    await deregister(pid)

async def handle_autopoll(identifier, args):
    log(f'handle_autopoll({identifier}, {args})')
    q = begin(identifier)
    end(q,identifier, 0)
    global AUTOPOLL
    if AUTOPOLL:
        return
    AUTOPOLL = 1

    global AUTOPOLL_TASK
    if AUTOPOLL_TASK is not None:
        return
    AUTOPOLL_TASK = asyncio.create_task(autopoll(1.0))

async def handle_poll(identifier, args):
    log(f'handle_poll({identifier}, {args})')
    output = await poll()
    log(f'handle_poll({identifier}, {args}): read {len(output)} categories of output')
    q = begin(identifier)
    if output is not None:
        send_poll_output(q,output)
        end(q,identifier, 0)
    else:
        end(q,identifier, 1)

async def handle_send(identifier, args):
    if len(args) < 2:
        fail(identifier, "not enough arguments")
        return
    try:
        pid = int(args[0])
        decoded = base64.b64decode(args[1])
    except Exception as e:
        log(f'handle_send: {e}')
        fail(identifier, "exception decoding argument")
        return
    if pid not in PROCESSES:
        log("No such process")
        q = begin(identifier)
        end(q,identifier, 1)
        return
    proc = PROCESSES[pid]
    log(f'write {decoded}')
    await proc.write(decoded)
    log('wrote')
    q = begin(identifier)
    end(q,identifier, 0)
    return False

async def handle_kill(identifier, args):
    log(f'kill {args}')
    try:
        pid = int(args[0])
    except:
        fail(identifier, "pid not an int")
        return
    if pid not in PROCESSES:
        log(f'no such process')
        q = begin(identifier)
        end(q,identifier, 1)
        return
    proc = PROCESSES[int(args[0])]
    proc.send_signal(signal.SIGTERM)
    q = begin(identifier)
    end(q,identifier, 0)
    return False

async def handle_quit(identifier, args):
    q = begin(identifier)
    end(q,identifier, 0)
    global AUTOPOLL_TASK
    if AUTOPOLL_TASK:
        log('will cancel autopoll')
        AUTOPOLL_TASK.cancel()
        try:
            log('await canceled autopoll task')
            await AUTOPOLL_TASK
        except asyncio.CancelledError:
            log('autopoll is now canceled')
        AUTOPOLL_TASK = None
    _it2_teardown_socket()
    return True

## Helpers for run()

async def start_process(args):
    runid = makeid()
    PROCESSES[runid] = proc
    return runid

def make_monitor_process(proc, islogin):
    def monitor_process(channel, value):
        log(f'monitor_process called with channel={channel} islogin={islogin} value={value}')
        if len(value) == 0:
            global COMPLETED
            log(f'add {proc.pid} to list of completed PIDs')
            COMPLETED.append(proc.pid)
            return cleanup()
        print_output(makeid(), proc.pid, channel, islogin, value)
        log("make_monitor_process: poll")
        try:
            poll_tty(proc)
        except Exception as e:
            log(f'make_monitor_process->poll_tty threw {e}: {traceback.format_exc()}')
        return None
    return monitor_process

def print_output(identifier, pid, channel, islogin, data):
    q = lock()
    if islogin:
        send_esc(q, f'%output {identifier} {pid} -1 {DEPTH}')
    else:
        send_esc(q, f'%output {identifier} {pid} {channel} {DEPTH}')
    send(q, data)
    send_esc(q, f'%end {identifier}')
    unlock(q)

def print_tty(message):
    q = lock()
    send_esc(q, f'%notif tty {message}')
    unlock(q)

## Infra

def fail(identifier, reason):
    log(f'fail: {reason}')
    q = begin(identifier)
    end(q, identifier, 1)

def begin(identifier):
    q = lock()
    send_esc(q, f'begin {identifier}')
    return q

def end(q, identifier, status):
    if len(PROCESSES):
        type = "f"
    else:
        type = "r"
    send_esc(q, f'end {identifier} {status} {type}')
    unlock(q)

async def cleanup():
    """Await tasks that have completed, clear the COMPLETED list, and remove them from TASKS."""
    log("cleaning up")
    global COMPLETED
    completed = list(COMPLETED)
    COMPLETED = []
    for pid in completed:
        if pid not in PROCESSES:
            log(f'pid {pid} no longer in PROCESSES, not cleaning up')
            continue
        log(f'clean up pid {pid}')
        proc = PROCESSES[pid]
        del PROCESSES[pid]
        await proc.cleanup()
        q = lock()
        send_esc(q, f'%terminate {proc.pid} {proc.return_code}')
        unlock(q)

async def handle(args):
    log(f'handle {args}')
    if len(args) == 0:
        # During recovery a blank line is sent. Just ignore it to avoid getting out of sync.
        return False
    cmd = args[0]
    del args[0]
    identifier = makeid()
    if cmd not in HANDLERS:
        fail(identifier, "unrecognized command")
        return

    f = HANDLERS[cmd]
    log(f'handler is {f}')
    should_quit = False
    try:
        should_quit = await f(identifier, args)
        if should_quit:
            global QUITTING
            QUITTING=True
    except Exception as e:
        log(f'Handler for {cmd} threw {e}: {traceback.format_exc()}')
    log("call cleanup()")
    await cleanup()

    global TASKS
    log(f'awaiting {TASKS}')
    while TASKS:
        task = TASKS[0]
        del TASKS[0]
        log(f'await {task}')
        await task
    TASKS=[]

    return should_quit

def read_line():
    try:
        log("Calling sys.stdin.readline")
        return sys.stdin.readline().rstrip('\n')
    except:
        log("Caught exception")
        sys.exit(1)

async def mainloop():
    global RUNLOOP
    global READSTATE
    RUNLOOP = asyncio.get_event_loop()
    args = []
    while True:
        log("reading")
        try:
            if READSTATE == 2:
                log('mainloop: send deferred ping')
                # It's safe to ping because we haven't sent a begin yet.
                ping()
            READSTATE = 1
            line = await asyncio.get_event_loop().run_in_executor(None, read_line)
        except:
            log("Exception while reading: " + traceback.format_exc())
            fail("none", "exception during read_line")
            return 0
        READSDTATE = 0
        log(f'read from stdin "{line}" with length {len(line)}')
        if len(line):
            if len(args) and args[-1].endswith("\\"):
                args[-1] = args[-1][:-1] + line
            else:
                args.append(line)
            log(f'args is now {args}')
        else:
            args = list(map(lambda s: base64.b64decode(s).decode('utf-8'), args))
            quit = await handle(args)
            log("flush stdout")
            sys.stdout.flush()
            if quit:
                log("Mainloop returns 0")
                return 0
            args = []

def ping():
    if len(PROCESSES):
        q = lock()
        send_esc(q, "%ping")
        unlock(q)
    else:
        log("Squelch pre-login ping")

async def update_pty_size():
    log(f'update_pty_size')
    window_size = fcntl.ioctl(sys.stdin.fileno(), termios.TIOCGWINSZ, '00000000')
    for pid in PROCESSES:
        proc = PROCESSES[pid]
        master = proc.master
        if master is not None:
            log(f'TIOCSWINSZ {proc}')
            # TODO: This is wrong becuse it could happen while awaiting something else between begin and end.
            fcntl.ioctl(master, termios.TIOCSWINSZ, window_size)
        else:
            log(f'no master fd for {proc}')

def on_sigwinch(_sig, _stack):
    log(f'Received SIGWINCH')
    if RUNLOOP is None:
        # There may not be an event loop yet.
        log('Ignore because no runloop')
        return
    global READSTATE
    if READSTATE == 1:
        log("send ping while blocked on stdin")
        ping()
    else:
        READSTATE = 2
    asyncio.run_coroutine_threadsafe(update_pty_size(), RUNLOOP)

# it2-over-ssh: proxy a remote unix socket to iTerm2. iTerm2's embedded it2 runs
# the command; framer.py is a dumb byte pipe. Frames up: "%it2 <connid> open",
# "%it2 <connid> data <base64>", "%it2 <connid> close". Commands down: "it2send
# <connid> <base64>", "it2close <connid>".
def _it2_teardown_socket():
    """Close the it2 server, unlink its socket, and release+remove its liveness lock.
    Idempotent and safe from atexit; everything lives in a persistent per-user dir so a
    leak would linger."""
    global IT2_SERVER, IT2_SOCKET_PATH, IT2_LOCK_FD, IT2_LOCK_PATH
    if IT2_SERVER is not None:
        with contextlib.suppress(Exception):
            IT2_SERVER.close()
        IT2_SERVER = None
    if IT2_SOCKET_PATH is not None:
        with contextlib.suppress(OSError):
            os.unlink(IT2_SOCKET_PATH)
        IT2_SOCKET_PATH = None
    if IT2_LOCK_PATH is not None:
        with contextlib.suppress(OSError):
            os.unlink(IT2_LOCK_PATH)
        IT2_LOCK_PATH = None
    if IT2_LOCK_FD is not None:
        with contextlib.suppress(OSError):
            os.close(IT2_LOCK_FD)  # closing the fd drops the advisory lock
        IT2_LOCK_FD = None

# Covers exits that skip handle_quit (read error, SIGHUP-then-clean-exit, etc.).
atexit.register(_it2_teardown_socket)

def _it2_close_connections():
    """Abort every accepted it2 client connection (but keep the listening socket).
    Called on recovery: iTerm2 restarted/reattached, so the demux that owned these
    connections is gone and their commands can never complete. Aborting each writer
    gives the orphaned it2.py client an EOF/RST so it exits instead of blocking forever
    on recv (its socket stayed open because this framer survived the restart)."""
    global IT2_CONNS
    for writer in list(IT2_CONNS.values()):
        # abort(), not close(): a graceful close() waits for the send buffer to flush
        # before sending FIN, which never completes for a stalled follow-stream whose
        # client stopped reading (exactly the mid-send connection this must unblock).
        # abort() RSTs immediately so the orphaned it2.py gets EOF/error and exits.
        with contextlib.suppress(Exception):
            writer.transport.abort()
    IT2_CONNS = {}

def _it2_lock_path(sock_path):
    return sock_path + ".lock"


def _it2_lock_is_held(lock_path):
    """True iff some live framer currently holds the exclusive advisory lock on lock_path.
    We test it non-blocking by trying to take it and dropping it immediately (closing the fd
    releases our test lock; a real owner's lock lives on a different fd and is untouched). A
    missing lock file is not held. An unopenable one is treated as held, i.e. we keep the
    socket rather than risk reclaiming a live one."""
    try:
        fd = os.open(lock_path, os.O_RDWR)
    except FileNotFoundError:
        return False
    except OSError:
        return True
    try:
        fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
        return False  # acquired -> nobody holds it
    except OSError:
        return True   # a live framer holds it
    finally:
        os.close(fd)


def _it2_sweep_stale_sockets(socket_dir, keep):
    """Reclaim it2 sockets leaked by prior sessions. When an ssh connection drops, framer is
    killed (SIGHUP/SIGKILL) before handle_quit or the atexit hook can unlink its socket, and
    each session used a fresh random name, so the owner-only dir accumulates dead sockets.
    Validity is decided by the liveness lock, not by connecting: the owner holds an exclusive
    flock on <socket>.lock for its whole life and the kernel drops it on death, so a socket
    whose lock is not held (dead owner) -- or has no lock at all (a leftover predating this
    scheme, e.g. the sockets already piled up on a host) -- is stale and reclaimable, while a
    held lock proves a live owner. Unlike a connect probe this cannot race a bind-in-progress
    or disturb a live peer. `keep` (the path we are about to bind) is skipped."""
    try:
        names = os.listdir(socket_dir)
    except OSError:
        return
    for name in names:
        full = os.path.join(socket_dir, name)
        if name.endswith(".sock") and full != keep:
            # Stale iff no live owner holds its lock (or there is no lock sidecar at all).
            if not _it2_lock_is_held(_it2_lock_path(full)):
                with contextlib.suppress(OSError):
                    os.unlink(full)
                with contextlib.suppress(OSError):
                    os.unlink(_it2_lock_path(full))
        elif name.endswith(".sock.lock") and full != _it2_lock_path(keep):
            # An orphan lock whose socket is already gone (owner died between locking and
            # binding, or after its socket was reclaimed above): drop it if unheld.
            sock_path = full[: -len(".lock")]
            if not os.path.exists(sock_path) and not _it2_lock_is_held(full):
                with contextlib.suppress(OSError):
                    os.unlink(full)


async def handle_it2listen(identifier, args):
    q = begin(identifier)
    global IT2_SERVER
    if IT2_SERVER is not None:
        end(q, identifier, 0)
        return
    if not args:
        end(q, identifier, 1)
        return
    path = args[0]
    global IT2_SOCKET_PATH, IT2_LOCK_FD, IT2_LOCK_PATH
    try:
        # Contain the socket in an owner-only (0700) directory. The socket file
        # itself is created 0700 (umask) then chmod 0600, but bind-time socket
        # permissions are not portable across platforms, so a 0700 parent is the
        # authoritative guard: no other local user can traverse to the socket to
        # connect to it, and none can pre-create/substitute it (this dir is under
        # $HOME, which only the owner can write). chmod fixes a pre-existing dir.
        socket_dir = os.path.dirname(path)
        if socket_dir:
            os.makedirs(socket_dir, mode=0o700, exist_ok=True)
            os.chmod(socket_dir, stat.S_IRWXU)
            # Reclaim sockets leaked by prior sessions killed before they could clean up.
            _it2_sweep_stale_sockets(socket_dir, keep=path)
        # Take the liveness lock BEFORE binding, so a concurrent sweeper always sees a live
        # owner for this socket and never reclaims it mid-bind. Held for our whole lifetime
        # via IT2_LOCK_FD; released by the kernel if we are killed. LOCK_NB fails only if
        # another live framer already owns this exact (random) path, which cannot happen in
        # practice -- treat it as a hard failure rather than blocking.
        lock_path = _it2_lock_path(path)
        lock_fd = os.open(lock_path, os.O_CREAT | os.O_RDWR, 0o600)
        try:
            fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except OSError:
            os.close(lock_fd)
            raise
        IT2_LOCK_FD = lock_fd
        IT2_LOCK_PATH = lock_path
        try:
            os.unlink(path)
        except OSError:
            pass
        # Bind under an owner-only umask so the socket is created as 0700 with no
        # window during which another local user could connect; the chmod then
        # narrows it to 0600 as belt-and-suspenders. umask only ever makes files
        # more restrictive, so leaking it to a concurrent task is harmless.
        old_umask = os.umask(0o077)
        try:
            IT2_SERVER = await asyncio.start_unix_server(handle_it2_connection, path=path)
        finally:
            os.umask(old_umask)
        os.chmod(path, stat.S_IRUSR | stat.S_IWUSR)
        IT2_SOCKET_PATH = path
        end(q, identifier, 0)
    except Exception as e:
        log(f'handle_it2listen failed: {e}')
        # Do not leave a live, possibly non-owner-restricted socket behind after
        # reporting failure: tear down whatever start_unix_server created and
        # reset the global so a later it2listen can retry instead of no-opping.
        if IT2_SERVER is not None:
            with contextlib.suppress(Exception):
                IT2_SERVER.close()
            IT2_SERVER = None
        IT2_SOCKET_PATH = None
        with contextlib.suppress(OSError):
            os.unlink(path)
        # Release + remove the liveness lock if we took it before the failure.
        if IT2_LOCK_FD is not None:
            with contextlib.suppress(OSError):
                os.close(IT2_LOCK_FD)
            IT2_LOCK_FD = None
        with contextlib.suppress(OSError):
            os.unlink(_it2_lock_path(path))
        IT2_LOCK_PATH = None
        end(q, identifier, 1)

async def handle_it2_connection(reader, writer):
    connid = makeid()
    # Register and emit `open` INSIDE the try, so that if the `open` emission (unlock's
    # os.write to the ssh up-channel) raises -- e.g. iTerm2's read side already closed --
    # the finally still pops IT2_CONNS and closes the writer instead of leaking both.
    try:
        IT2_CONNS[connid] = writer
        q = lock()
        send_esc(q, f'%it2 {connid} open')
        unlock(q)
        while True:
            data = await reader.read(IT2_READ_CHUNK)
            if not data:
                break
            q = lock()
            send_esc(q, f'%it2 {connid} data {base64.b64encode(data).decode("ascii")}')
            unlock(q)
    except Exception as e:
        log(f'it2 connection {connid}: {e}')
    finally:
        IT2_CONNS.pop(connid, None)
        # Emit `close` UNLESS iTerm2 itself initiated the teardown via it2close (then it already
        # forgot the connid and a close echo would be a duplicate/unsolicited frame). A
        # framer/client-initiated teardown -- client disconnect, or a buffer-cap abort that
        # popped connid -- is NOT flagged here, so it still emits close so iTerm2 learns the
        # connection died. Suppress() so a failing close emission cannot skip writer.close().
        locally_closed = connid in IT2_LOCALLY_CLOSED
        IT2_LOCALLY_CLOSED.discard(connid)
        if not locally_closed:
            with contextlib.suppress(Exception):
                q = lock()
                send_esc(q, f'%it2 {connid} close')
                unlock(q)
        with contextlib.suppress(Exception):
            writer.close()

async def handle_it2send(identifier, args):
    q = begin(identifier)
    if len(args) < 2 or args[0] not in IT2_CONNS:
        end(q, identifier, 1)
        return
    writer = IT2_CONNS[args[0]]
    # Sample the unsent buffer BEFORE writing. The cap guards against a *stalled* consumer
    # whose buffer grows without bound across successive sends; a single frame onto a healthy
    # (near-empty) buffer is bounded memory, not a leak, so it must pass even if it alone
    # exceeds the cap. Sampling AFTER the write would abort a perfectly healthy fast client on
    # one oversized down-frame: write() does a single optimistic synchronous send (bounded by
    # the ~200 KB socket wmem) and buffers the rest, and there is no await before the sample,
    # so the event loop never runs to flush first -> the whole frame reads as over-cap. A
    # genuinely wedged consumer is still caught on its NEXT send, when the accumulated unsent
    # buffer is over the cap before we even write.
    try:
        buffered = writer.transport.get_write_buffer_size()
    except Exception:
        buffered = 0
    if buffered > IT2_WRITE_BUFFER_MAX:
        log(f'handle_it2send: write buffer {buffered} over cap for {args[0]}; dropping connection')
        IT2_CONNS.pop(args[0], None)
        with contextlib.suppress(Exception):
            writer.transport.abort()
        end(q, identifier, 1)
        return
    # Decode OUTSIDE the socket try/except: a malformed base64 payload (bad padding, a
    # corrupted/truncated command line) raises binascii.Error/ValueError, which is NOT a dead
    # socket. Skip just this frame and keep the connection alive rather than tearing it down
    # (pop+abort) the way a genuine write failure must.
    try:
        payload = base64.b64decode(args[1])
    except (ValueError, TypeError) as e:
        log(f'handle_it2send: bad base64 for {args[0]}: {e}; dropping this frame')
        end(q, identifier, 1)
        return
    try:
        writer.write(payload)
    except Exception as e:
        # A write failure (e.g. the peer RST'd the socket mid-stream ->
        # ConnectionResetError/BrokenPipeError) leaves a dead writer. Drop and abort it so
        # later sends do not re-enter and repeat the failing write.
        log(f'handle_it2send: {e}')
        IT2_CONNS.pop(args[0], None)
        with contextlib.suppress(Exception):
            writer.transport.abort()
        end(q, identifier, 1)
        return
    # Do NOT await drain() here. This is the single serial command loop (mainloop does
    # `await handle(args)` and reads no further command until it returns), so blocking on a
    # stalled consumer's backpressure would freeze send/kill/file/poll for the WHOLE ssh
    # session. write() only buffers; asyncio flushes it in the background as the socket
    # drains, so a live-but-slow consumer keeps working.
    end(q, identifier, 0)

async def handle_it2close(identifier, args):
    q = begin(identifier)
    connid = args[0] if args else None
    writer = IT2_CONNS.pop(connid, None) if connid is not None else None
    if writer is not None:
        # Flag this as an iTerm2-initiated close so the still-live reader coroutine's finally
        # does not echo a duplicate close back up. Only meaningful while the reader is alive
        # (connid was still registered); it discards the flag when it unwinds.
        IT2_LOCALLY_CLOSED.add(connid)
        # abort(), not close(): a graceful close() defers FIN until the send buffer flushes,
        # which never happens for a stalled follow-stream client (paused stdout consumer, so the
        # down buffer is non-empty). That would leave the socket fd open, the reader coroutine
        # blocked on read() forever, and its IT2_LOCALLY_CLOSED flag leaked for the ssh
        # session's life. abort() RSTs immediately, giving the reader EOF so it unwinds. Matches
        # the twin _it2_close_connections().
        with contextlib.suppress(Exception):
            writer.transport.abort()
    end(q, identifier, 0)


HANDLERS = {
    "run": handle_run,
    "login": handle_login,
    "send": handle_send,
    "kill": handle_kill,
    "quit": handle_quit,
    "register": handle_register,
    "deregister": handle_deregister,
    "poll": handle_poll,
    "reset": handle_reset1,
    "reset2": handle_reset2,
    "autopoll": handle_autopoll,
    "recover": handle_recover,
    "save": handle_save,
    "file": handle_file,
    "eval": handle_eval,
    "getenv": handle_getenv,
    "runpy": handle_runpy,
    "it2listen": handle_it2listen,
    "it2send": handle_it2send,
    "it2close": handle_it2close
}

def _silence_shutdown_unraisable(unraisable):
    # Subprocess transports' __del__ runs after asyncio.run() has closed the
    # loop; their cleanup calls call_soon() and raises "Event loop is closed".
    # We're on our way out — nobody cares.
    exc = unraisable.exc_value
    if isinstance(exc, RuntimeError) and 'Event loop is closed' in str(exc):
        return
    sys.__unraisablehook__(unraisable)

def main():
    if sys.stdin.isatty():
        signal.signal(signal.SIGWINCH, on_sigwinch)
    sys.unraisablehook = _silence_shutdown_unraisable
    try:
        asyncio.run(mainloop())
    except Exception as e:
        log(f'Exception {traceback.format_exc()}')

if __name__ == "__main__":
    main()
