"""Run CREDIT rollouts from the notebook in a subprocess, with live progress output."""

import codecs
import fcntl
import os
import pty
import select
import struct
import subprocess
import sys
import termios
from pathlib import Path

LAUNCHER = Path(__file__).with_name("credit_launch.py")


def run_configs(configs, fast_math=False):
    """Run rollout configs one after another in ONE subprocess (the model is loaded once; see credit_launch.py).

    The subprocess uses this kernel's interpreter and gets the kernel's sys.path as PYTHONPATH, so it finds packages
    the kernel finds (e.g. a `pip install --target` directory added to sys.path). Output goes through a pseudo-terminal
    so CREDIT's tqdm bars refresh in place instead of printing a line per update.
    """
    configs = [str(c) for c in configs]
    cmd = [sys.executable, str(LAUNCHER), *(["--fast"] if fast_math else []), *configs]
    print(f"$ credit rollout x {len(configs)} run(s){' (FAST_MATH)' if fast_math else ''}", flush=True)
    env = {**os.environ, "PYTHONPATH": os.pathsep.join(p for p in sys.path if p)}

    master, slave = pty.openpty()
    # A fresh pty reports 0 columns, and tqdm then draws no bar; give it a normal terminal size.
    fcntl.ioctl(slave, termios.TIOCSWINSZ, struct.pack("HHHH", 24, 120, 0, 0))
    proc = subprocess.Popen(cmd, stdout=slave, stderr=slave, close_fds=True, env=env)
    os.close(slave)  # the parent only reads the master end
    decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")  # reads can split multi-byte characters

    def forward(timeout):
        # Copy whatever the child wrote to the notebook. False when nothing was read.
        if not select.select([master], [], [], timeout)[0]:
            return False
        try:
            data = os.read(master, 1024)
        except OSError:  # Linux raises EIO once the child has exited and the pty is empty
            return False
        if not data:
            return False
        sys.stdout.write(decoder.decode(data))
        sys.stdout.flush()
        return True

    try:
        while True:
            if forward(0.1):
                continue
            if proc.poll() is not None:
                while forward(0):  # drain output still buffered after the child exited
                    pass
                break
    finally:
        os.close(master)
        proc.wait()
    if proc.returncode != 0:
        raise RuntimeError(f"credit rollout exited with code {proc.returncode}")