import asyncio
import os
import sys
import signal
import subprocess

async def _kill_process_tree(pid: int):
    if sys.platform != "win32":
        try:
            pgid = os.getpgid(pid)
            os.killpg(pgid, signal.SIGTERM)
            for _ in range(30):
                await asyncio.sleep(0.1)
                try:
                    os.killpg(pgid, 0)
                except ProcessLookupError:
                    return
            os.killpg(pgid, signal.SIGKILL)
        except ProcessLookupError:
            pass
        except Exception as e:
            print(f"Error killing process group {pid}: {e}")
    else:
        try:
            subprocess.run(["taskkill", "/F", "/T", "/PID", str(pid)], capture_output=True)
        except Exception as e:
            print(f"Error killing Windows process tree {pid}: {e}")

def pid_exists(pid: int) -> bool:
    try:
        os.kill(pid, 0)
        return True
    except OSError:
        return False

def get_child_pids(parent_pid: int) -> list[int]:
    if sys.platform == "win32":
        return []
    try:
        res = subprocess.run(["pgrep", "-P", str(parent_pid)], capture_output=True, text=True)
        if res.returncode == 0:
            return [int(p.strip()) for p in res.stdout.splitlines() if p.strip()]
    except Exception:
        pass
    return []

async def main():
    # Use sys.executable to ensure we run with the correct python interpreter
    script = (
        f"import subprocess, time; "
        f"p = subprocess.Popen(['{sys.executable}', '-c', 'import time; time.sleep(100)']); "
        f"time.sleep(100)"
    )
    
    kwargs = {}
    if sys.platform != "win32":
        kwargs["preexec_fn"] = os.setpgrp
        
    proc = await asyncio.create_subprocess_exec(
        sys.executable, "-c", script,
        **kwargs
    )
    
    pid = proc.pid
    print(f"Spawned parent process with PID {pid}")
    
    await asyncio.sleep(1.5)
    
    children = get_child_pids(pid)
    print(f"Found children: {children}")
    if sys.platform != "win32" and not children:
        print("RED Test failed: No child processes found.")
        sys.exit(1)
        
    child_pid = children[0] if children else None
    
    print(f"Killing process tree for PID {pid}...")
    await _kill_process_tree(pid)
    
    await asyncio.sleep(1.0)
    
    parent_alive = pid_exists(pid)
    child_alive = pid_exists(child_pid) if child_pid else False
    
    print(f"Parent alive: {parent_alive}, Child alive: {child_alive}")
    if parent_alive or child_alive:
        print("RED Test failed: Process or child is still alive!")
        sys.exit(1)
    else:
        print("GREEN Test PASSED: Both parent and child processes successfully terminated via process tree cleanup!")
        sys.exit(0)

if __name__ == "__main__":
    asyncio.run(main())