#!/usr/bin/env python3
"""PTY helper: creates a real pseudo-terminal and proxies stdin/stdout.
Usage: python3 pty-helper.py <shell> <cwd> [cols] [rows]

If the AGENTSROOM_INIT_CMD env var is set, it will be sent to the shell
after initialization (once the shell produces its first output).

When the parent opens a 4th pipe (fd 3), the helper also reports foreground-job
transitions on it as JSON lines (see foreground_watcher below). fd 1/2 stay the
raw terminal stream, so the control channel can never pollute the scrollback.
"""
import pty, os, sys, struct, fcntl, termios, signal, time, tempfile, threading, subprocess, json

shell = sys.argv[1] if len(sys.argv) > 1 else '/bin/zsh'
cwd = sys.argv[2] if len(sys.argv) > 2 else os.path.expanduser('~')
cols = int(sys.argv[3]) if len(sys.argv) > 3 else 120
rows = int(sys.argv[4]) if len(sys.argv) > 4 else 30

# Capture startup & init commands before fork
startup_cmd = os.environ.pop('AGENTSROOM_STARTUP_CMD', '')
init_cmd = os.environ.pop('AGENTSROOM_INIT_CMD', '')

# Write long commands to a temp script to avoid PTY buffer limits (~4096 bytes)
init_script_path = None
if init_cmd:
    fd_tmp, init_script_path = tempfile.mkstemp(prefix='agentsroom-init-', suffix='.sh')
    os.write(fd_tmp, init_cmd.encode('utf-8'))
    os.write(fd_tmp, b'\n')
    os.close(fd_tmp)

# Fork with PTY
pid, fd = pty.fork()

if pid == 0:
    # Child: set up and exec shell
    os.chdir(cwd)
    os.execvp(shell, [shell, '-l'])
    sys.exit(1)

# Parent: set initial window size
try:
    winsize = struct.pack('HHHH', rows, cols, 0, 0)
    fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize)
except Exception:
    pass

# Graceful shutdown: kill every process in the shell's session.
# pty.fork() makes the child a session leader (sid == shell pid). Interactive
# shells use job control, so each foreground command (yarn dev, vite, ...)
# runs in its OWN process group — killpg on the shell's pgrp would leave them
# behind as ghosts. We walk the whole session instead.
def kill_session(sig):
    delivered = False
    try:
        out = subprocess.check_output(
            ['pgrep', '-s', str(pid)],
            stderr=subprocess.DEVNULL,
            timeout=2,
        ).decode('utf-8', errors='ignore')
        for token in out.split():
            try:
                target = int(token)
            except ValueError:
                continue
            if target <= 0 or target == os.getpid():
                continue
            try:
                os.kill(target, sig)
                delivered = True
            except (ProcessLookupError, PermissionError):
                pass
    except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired):
        pass
    if delivered:
        return
    # Fallback: kill the shell's process group, then the shell itself.
    try:
        os.killpg(os.getpgid(pid), sig)
    except Exception:
        try:
            os.kill(pid, sig)
        except Exception:
            pass

def handle_sigterm(sig, frame):
    kill_session(signal.SIGTERM)
    time.sleep(0.5)
    kill_session(signal.SIGKILL)
    sys.exit(0)

signal.signal(signal.SIGTERM, handle_sigterm)


def write_all(target_fd, data):
    """Write ALL bytes to fd, looping on short writes."""
    mv = memoryview(data)
    while len(mv) > 0:
        try:
            n = os.write(target_fd, mv)
            mv = mv[n:]
        except OSError:
            return False
    return True


# --- Thread: PTY → stdout (sends terminal output to Electron) ---
def pty_reader():
    """Continuously read from PTY master and write to stdout."""
    stdout_fd = 1  # fd 1 = stdout, always
    try:
        while True:
            try:
                data = os.read(fd, 65536)
            except OSError:
                break
            if not data:
                break
            if not write_all(stdout_fd, data):
                break
    except Exception:
        pass


# --- Thread: stdin → PTY (sends user input to the shell) ---
def stdin_reader():
    """Continuously read from stdin and write to PTY master."""
    stdin_fd = 0  # fd 0 = stdin, always
    try:
        while True:
            try:
                data = os.read(stdin_fd, 65536)
            except OSError:
                break
            if not data:
                break

            # Check for resize escape sequence: \x1b_RESIZE;rows;cols\x1b\\
            text = data
            while b'\x1b_RESIZE;' in text:
                idx = text.index(b'\x1b_RESIZE;')
                # Write everything before the escape to PTY
                if idx > 0:
                    write_all(fd, text[:idx])
                end = text.find(b'\x1b\\', idx)
                if end == -1:
                    # Incomplete escape — discard it to avoid corrupting PTY input
                    text = b''
                    break
                resize_data = text[idx + 9:end].decode('utf-8', errors='ignore')
                parts = resize_data.split(';')
                if len(parts) == 2:
                    try:
                        r_rows, r_cols = int(parts[0]), int(parts[1])
                        winsize = struct.pack('HHHH', r_rows, r_cols, 0, 0)
                        fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize)
                    except Exception:
                        pass
                text = text[end + 2:]

            # Write remaining normal data to PTY
            if text:
                write_all(fd, text)
    except Exception:
        pass


# --- Thread: foreground job watcher (control channel on fd 3) ---
#
# Tells "a command is running" apart from "the shell is back at its prompt",
# without parsing terminal output and without injecting anything into the
# user's shell configuration.
#
# pty.fork() makes the shell a session leader, so the shell's process group id
# equals its pid. Interactive shells run each foreground job in its OWN process
# group, so the terminal's foreground pgrp (TIOCGPGRP, exposed as
# os.tcgetpgrp on the master fd) is:
#   == shell pid  -> the shell is at its prompt (idle)
#   != shell pid  -> a foreground job is running (busy), and the value IS that
#                    job's pgid, which the parent can resolve to a command line
#   <= 0          -> the session is gone
#
# One syscall per tick, so the poll is free. The parent decides what to do with
# the transitions (see electron/ipc/ptyForeground.ts).
CONTROL_FD = 3
FG_POLL_INTERVAL_S = 1.0


def control_channel_available():
    """True when the parent actually opened fd 3 for us."""
    try:
        os.fstat(CONTROL_FD)
        return True
    except OSError:
        return False


def foreground_watcher():
    last_state = None
    while True:
        time.sleep(FG_POLL_INTERVAL_S)
        try:
            fg = os.tcgetpgrp(fd)
        except OSError:
            break
        if fg <= 0:
            break  # session torn down; the exit path reports it
        state = 'idle' if fg == pid else 'busy'
        # Report a NEW busy job even when the previous state was also busy: two
        # chained commands (`a ; b`) get distinct process groups with no idle
        # tick in between at this cadence.
        key = (state, fg)
        if key == last_state:
            continue
        last_state = key
        payload = {'type': state}
        if state == 'busy':
            payload['pgid'] = fg
        try:
            os.write(CONTROL_FD, (json.dumps(payload) + '\n').encode('utf-8'))
        except OSError:
            break  # parent closed the channel; stop watching, keep the terminal alive


# Start I/O threads
t_pty = threading.Thread(target=pty_reader, daemon=True)
t_stdin = threading.Thread(target=stdin_reader, daemon=True)
t_pty.start()
t_stdin.start()

if control_channel_available():
    threading.Thread(target=foreground_watcher, daemon=True).start()

# Send startup command (e.g. "bash") then init command once the shell is ready
if startup_cmd or init_script_path:
    # Wait a bit for the shell to initialize
    time.sleep(0.5)
    try:
        if startup_cmd:
            # Run startup command first (e.g. switch shell, source a file)
            write_all(fd, (' ' + startup_cmd + '\n').encode('utf-8'))
            # Wait for the startup command to settle before sending init
            time.sleep(0.8)
        if init_script_path:
            # Leading space prevents zsh from recording this in history
            # (requires HIST_IGNORE_SPACE which is default in oh-my-zsh).
            # 'clear' wipes the shell profile noise (aliases, motd, etc.) before
            # the init command runs so the terminal starts clean.
            short_cmd = (' clear ; source ' + init_script_path + ' ; rm -f ' + init_script_path + '\n')
            write_all(fd, short_cmd.encode('utf-8'))
    except Exception:
        pass

# Wait for PTY reader to finish (child process exited)
t_pty.join()

# Cleanup
if init_script_path:
    try:
        os.unlink(init_script_path)
    except Exception:
        pass

try:
    os.close(fd)
except Exception:
    pass

try:
    _, status = os.waitpid(pid, 0)
    sys.exit(os.WEXITSTATUS(status) if os.WIFEXITED(status) else 1)
except Exception:
    sys.exit(0)
