import json
import os
import select
import socket
import subprocess


def handle_run(run: dict) -> dict:
    res = {}
    text = run.get("text", False)
    env = {
        **os.environ,
        "PATH": ":".join(
            os.environ.get("PATH", "").split(":") + run.get("EXTRA_PATH", [])
        ),
    }
    proc = None
    try:
        proc = subprocess.Popen(
            req["run"]["argv"],
            text=text,
            env=env,
            cwd="/home/user",
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
        )
        res["status"] = "exec succeeded"
    except Exception as e:
        res["status"] = "exec failed"
        res["exception"] = repr(e)
    res["pid"] = getattr(proc, "pid", None)
    try:
        if proc is not None:
            proc.wait(0.125)
        res["long_running"] = False
        res["returncode"] = getattr(proc, "returncode", None)
    except subprocess.TimeoutExpired:
        res["long_running"] = True
    return res, proc


if __name__ == "__main__":
    serv = socket.fromfd(3, socket.AF_VSOCK, socket.SOCK_STREAM)

    procs = []
    conns = [serv]

    while True:
        rr, rw, xs = select.select(conns, [], [])

        for con in rr:
            if con is serv:
                con, (cid, port) = serv.accept()
                assert cid == 2, cid
                conns.append(con)
                continue
            req = con.recv(8192)
            # IDK why but I keep getting empty messages
            if req == b"":
                continue
            try:
                req = json.loads(req)
                print(f"Received {req=}")
            except json.JSONDecodeError as e:
                print(f"Couldn't interpret {req=}: {e}")
                continue
            if "run" in req:
                res, proc = handle_run(req["run"])
                procs.append(proc)
            else:
                res = {"status": "unknown command"}
            _, rw, _ = select.select([], [con], [])
            assert rw, rw
            res = json.dumps(res).encode("utf8")
            print(f"Responding with {res=}")
            con.send(res)