# 检测 python 脚本是否运行，一旦检测到程序停止就启动脚本并且将输出显示到当前屏幕上
# 按下 ctrl + c 后(按一次就行)会依次对接管的脚本发送 ctrl + c，使程序能捕获 KeyboardInterrupt 异常并退出

import psutil
import os
import sys
import pexpect
from loguru import logger
from time import sleep
from concurrent.futures import ThreadPoolExecutor


COMMANDS = [('python3.10', 'u2_magic.py'),  # 可以填多个
            # ('python3.10', 'auto_magic_seed.py'),
            ]



class Monitor:
    def __init__(self, interpreter, script):
        self.interpreter = interpreter
        self.script = script
        self.child = None

    def run(self):
        while True:
            try:
                i = 0
                for pid in psutil.pids():
                    try:
                        process = psutil.Process(pid)
                        if self.interpreter in process.name():
                            if any(self.script in s for s in process.cmdline()):
                                i = 1
                    except psutil.NoSuchProcess:
                        pass
                    except Exception as e:
                        logger.exception(e)
                if i == 0:
                    logger.error(f'Process {self.script} is down')
                    if self.child:
                        self.child.terminate()
                    self.child = pexpect.spawn(f'{self.interpreter} {self.script}', logfile=sys.stdout.buffer)
                    logger.warning(f'Run command | {self.interpreter} {self.script}')
                    try:
                        self.child.expect(pexpect.EOF, timeout=9223372036)
                    except Exception as e:
                        logger.exception(e)
                elif self.child:
                    try:
                        self.child.expect(pexpect.EOF, timeout=9223372036)
                    except Exception as e:
                        logger.exception(e)
            except Exception as e:
                logger.exception(e)
            finally:
                sleep(10)


if __name__ == '__main__':
    log_path = f'{os.path.splitext(__file__)[0]}.log'
    logger.add(level='DEBUG', sink=log_path, encoding='utf-8', rotation="5 MB")
    monitors = [Monitor(interpreter, script) for interpreter, script in COMMANDS]
    try:
        with ThreadPoolExecutor(max_workers=len(monitors)) as executor:
            [executor.submit(monitor.run) for monitor in monitors]
    except KeyboardInterrupt:
        for monitor in monitors:
            if monitor.child:
                monitor.child.sendcontrol('c')
                sleep(0.1)
                monitor.child.sendcontrol('c')
                logger.info(f'Sent Ctrl + C to the process {monitor.script}')
            else:
                logger.info(f'Process {monitor.script} is still running')
        sys.exit()