import sys
import psutil

from typing import Optional


def get_pid_of_current_script() -> int:
    for proc in psutil.process_iter():
        try:
            proc_cmd = proc.cmdline()
            if proc_cmd[1:] == sys.argv:
                return proc.pid
        except (psutil.AccessDenied, psutil.ZombieProcess):
            pass
    raise RuntimeError("failed to find pid")


def get_pid_of_python_script(script_name: str) -> Optional[int]:
    """
    This is very brittle, don't expect it to always work. 
    TODO (swapnil) - generalize this a bit better. This was only ever tested on Ubuntu
    """
    for proc in psutil.process_iter():
        try:
            proc_cmd = proc.cmdline()
            if len(proc_cmd) >= 2 and "python3" in proc_cmd[0] and script_name in proc_cmd[1]:
                return proc.pid
        except psutil.AccessDenied:
            pass
    return None
