''' opencos.utils.subprocess_helpers -- wrappers for subprocess to support background/tee'''

import os
import signal
import subprocess
import sys
import threading

import psutil

from opencos import util
from opencos.files import safe_shutil_which
from opencos.util import debug, error, info, warning, progname, global_log
from opencos.utils import status_constants
from opencos.utils.str_helpers import strip_ansi_color

IS_LINUX: bool = sys.platform.startswith('linux')
IS_MACOS: bool = sys.platform.startswith('darwin')
IS_WINDOWS: bool = sys.platform.startswith('win')

# For non-Windows, we track the background parent PIDs, because some tools (vivado XSim,
# most Modelsim/Questa variants) tend to spawn children PIDs that don't always respond
# nicely to a friendly *nix SIGTERM. So we'll remember what our parent PIDs are, and
# eda.py's (or other CLI opencos script) can use signal to cleanup any remaining
# parents + children using subprocess_helpers.cleanup_all()
#
# Note that in non-Windows we run all subprocess with start_new_session=True, to create
# chain them to a PGID group for more effective terminate / killing later. In Windows
# jobs we use: creationflags=subprocess.CREATE_NEW_PROCESS_GROUP
ALL_PARENT_PIDS = set()

def get_command_list(command_list: list, shell: bool = False):
    '''Returns a str if this is non-Windows + shell=True'''
    if not IS_WINDOWS and shell:
        # Return a str:
        return ' '.join(command_list)

    return command_list # leave as list.


def get_kwargs(
        work_dir: str, shell: bool = False, is_subprocess_popen: bool = True
) -> dict:
    '''Returns the subprocess.Popen class ref'''
    proc_kwargs = { 'shell': shell }

    if is_subprocess_popen:
        proc_kwargs.update({
            'stdout': subprocess.PIPE,
            'stderr': subprocess.STDOUT
        })

    if work_dir:
        proc_kwargs['cwd'] = work_dir

    bash_exec = safe_shutil_which('bash')
    if shell and bash_exec and not IS_WINDOWS:
        # Note - windows powershell will end up calling: /bin/bash /c, which won't work
        proc_kwargs['executable'] = bash_exec    # for bash features (why we're running shell=True)

    if IS_WINDOWS:
        proc_kwargs['creationflags'] = subprocess.CREATE_NEW_PROCESS_GROUP
    else:
        proc_kwargs['start_new_session'] = True

    return proc_kwargs



def subprocess_run(
        work_dir: str, command_list: list, fake: bool = False, shell: bool = False,
        timeout: int = 0
) -> int:
    ''' Run command_list in the foreground, with preference to use bash if shell=True'''

    proc_kwargs = get_kwargs(work_dir=work_dir, shell=shell, is_subprocess_popen=False)

    # For subprocess.run(...) with timeout we simply add to args,
    # pass timeout=value to subprocess.run(..)
    if timeout > 0:
        proc_kwargs['timeout'] = timeout

    c = get_command_list(command_list, shell=shell)

    if fake:
        info(f"subprocess_run FAKE: would have called subprocess.run({c}, **{proc_kwargs}")
        return 0

    debug(f"subprocess_run: About to call subprocess.run({c}, **{proc_kwargs}")
    proc = subprocess.run(c, check=True, **proc_kwargs)
    # Note - we do not get PID management for subprocess_run(...)
    return proc.returncode


# pylint: disable=unknown-option-value
# pylint: disable=too-many-branches,too-many-locals,too-many-statements,too-many-arguments
# pylint: disable=too-many-positional-arguments
def subprocess_run_background(
        work_dir: str, command_list: list, background: bool = True, fake : bool = False,
        shell: bool = False, tee_fpath: str = '',
        timeout: int = 0, timeout_str: str = '',
        progress_timeout: int = 0, progress_timeout_str: str = '',
        force_no_color: bool = False
) -> (str, str, int):
    ''' Run command_list in the background, with preference to use bash if shell=True

    tee_fpath is relative to work_dir.

    Note that stderr is converted to stdout, and stderr is retuned as '':
        Returns tuple of (stdout str, '', int return code)
    '''

    debug(f'subprocess_run_background: {background=} {tee_fpath=} {shell=}')

    if fake:
        # let subprocess_run handle it (won't run anything)
        rc = subprocess_run(work_dir, command_list, fake=fake, shell=shell, timeout=timeout)
        return '', '', rc

    proc_kwargs = get_kwargs(work_dir=work_dir, shell=shell, is_subprocess_popen=True)

    c = get_command_list(command_list, shell=shell)

    debug(f"subprocess_run_background: about to call subprocess.Popen({c}, **{proc_kwargs})")

    # Set-up a watchdog to track timeout, if timeout > 0 (enabled)
    def overall_watchdog_timer_expired(proc) -> None:
        error(f'subprocess_run_background: TimeoutExpired after {timeout} seconds {timeout_str}',
              do_exit=False, error_code=status_constants.EDA_EXEC_SUBPROCESS_TIMEOUT)
        kill_proc_tree(proc.pid)

    # Set-up another watchdog to track progress timeout, if timeout > 0 (enabled)
    def progress_watchdog_timer_expired(proc) -> None:
        error(f'subprocess_run_background: TimeoutExpired after {progress_timeout} seconds',
              f'{progress_timeout_str}',
              do_exit=False, error_code=status_constants.EDA_EXEC_SUBPROCESS_TIMEOUT)
        kill_proc_tree(proc.pid)

    proc = subprocess.Popen(c, **proc_kwargs) # pylint: disable=consider-using-with
    if not background:
        info(f'PID {proc.pid} for {command_list[0]}')
    add_running_parent_pid(proc.pid)

    # This will fire even if readline() is blocked
    if timeout > 0 and proc.pid:
        watchdog = threading.Timer(timeout, overall_watchdog_timer_expired, args=[proc])
        watchdog.start()
    else:
        watchdog = None

    progress_watchdog = None

    def reset_progress_watchdog_timer():
        nonlocal progress_watchdog

        if progress_timeout == 0 or not proc.pid:
            # not enabled, don't start it.
            return

        if progress_watchdog:
            progress_watchdog.cancel()

        # Reuse your existing expire logic
        progress_watchdog = threading.Timer(
            progress_timeout, progress_watchdog_timer_expired,  args=[proc]
        )
        progress_watchdog.start()


    stdout_chunks = []
    tee_fpath_f = None
    reset_progress_watchdog_timer()

    if tee_fpath:
        tee_fpath = os.path.join(work_dir, tee_fpath)
        try:
            tee_fpath_f = open( # pylint: disable=consider-using-with
                tee_fpath, 'w', encoding='utf-8'
            )
        except Exception as e:
            error(f'Unable to open file "{tee_fpath}" for writing, {e}')

    no_color = force_no_color or not util.args['color']
    try:
        for line in iter(proc.stdout.readline, b''):
            line = line.decode("utf-8", errors="replace") # leave \n intact

            reset_progress_watchdog_timer()

            # Since we don't control what the subprocess command did, if it
            # thinks we support color, but user ran with --no-color, we need to strip ANSI colors:
            if no_color:
                line = strip_ansi_color(line)

            # Print the line with color, if --color:
            if not background:
                print(line, end='')

            # for all logs, and the returned stdout str, if we haven't stripped color yet,
            # we need to now, before writing to tee_fpath_f, or to global_log:
            if not no_color:
                line = strip_ansi_color(line)
            line = line.replace('\r', '') # remove CR

            if tee_fpath_f:
                tee_fpath_f.write(line)
            if global_log.file:
                # directly write to file handle, avoid util.UtilLogger.write(line, end='')
                global_log.file.write(line)
            stdout_chunks.append(line)

        # If we broke the for-loop, then stdout has expired and we broke the process.
        proc.communicate()

    finally:
        if watchdog:
            watchdog.cancel()
        if progress_watchdog:
            progress_watchdog.cancel()

    remove_completed_parent_pid(proc.pid)

    rc = proc.returncode
    if tee_fpath_f:
        tee_fpath_f.write(f'INFO: [{progname}] subprocess_run_background: returncode={rc}\n')
        tee_fpath_f.close()
        if not background:
            info('subprocess_run_background: wrote: ' + os.path.abspath(tee_fpath))

    return ''.join(stdout_chunks), '', rc


def add_running_parent_pid(pid: int) -> None:
    '''Adds pid (if still alive) to ALL_PARENT_PIDS'''
    try:
        p = psutil.Process(pid)
        ALL_PARENT_PIDS.add(p.pid)
    except (ProcessLookupError, psutil.NoSuchProcess):
        pass
    except Exception as e:
        error(f'{pid=} exception {e}')

def remove_completed_parent_pid(pid: int) -> None:
    '''Removes pid (if no longer alive) from ALL_PARENT_PIDS.'''
    try:
        p = psutil.Process(pid)
        warning(f'PID {p.pid} still running')
    except (ProcessLookupError, psutil.NoSuchProcess):
        ALL_PARENT_PIDS.remove(pid)
    except Exception as e:
        error(f'{pid=} exception {e}')


def cleanup_all() -> None:
    '''Kills everything from ALL_PARENT_PIDS.'''
    for parent in ALL_PARENT_PIDS:
        kill_proc_tree(parent)


def kill_proc_tree(pid: int) -> None:
    '''
    Kills a process and its entire descendant tree

    Will attempt to use (non-Windows) the process-groupid (pgid)
    '''

    # For Windows:
    #
    # Send CTRL_BREAK to the entire group
    # Note: This requires CREATE_NEW_PROCESS_GROUP at startup
    #os.kill(proc.pid, signal.CTRL_BREAK_EVENT)
    #
    # /F = Force
    # /T = Tree (Kills the process AND all child processes)
    # /PID = The process ID
    #subprocess.run(['taskkill', '/F', '/T', '/PID', str(proc.pid)],
    #               stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


    # 1. Grab PGID immediately if we can. This is to deal with subprocesses
    # shell=True, or closed sourced tools that spawned their own ways:
    pgid = None
    if not IS_WINDOWS:
        try:
            pgid = os.getpgid(pid) # pylint: disable=no-member
        except (ProcessLookupError, OSError):
            pgid = None

    # 2. Get all processes to kill
    all_to_kill = {} # { key=pid: value=psutil.Process }
    try:
        parent = psutil.Process(pid)
        for p in parent.children(recursive=True) + [parent]:
            all_to_kill[p.pid] = p
        info(f'{pid=} {parent=}, {all_to_kill=}')
    except (ProcessLookupError, psutil.NoSuchProcess):
        all_to_kill = {}

    # 2b. Additional fallback, search the system for anyone else in the same
    # PGID to escape processes that above parent.childen(recursive=True) may have
    # missed
    if pgid:
        for p in psutil.process_iter(['pid']):
            if not p:
                continue
            try:
                # Manually fetch the PGID using the standard os module
                # pylint: disable=no-member
                if os.getpgid(p.info['pid']) == pgid and p.info['pid'] not in all_to_kill:
                    all_to_kill[p.info['pid']] = p
            except (ProcessLookupError, psutil.NoSuchProcess, psutil.AccessDenied):
                continue

    # 3. Start gracefull killing (terminate / SIGTERM):
    for p in all_to_kill.values():
        if not p:
            continue
        try:
            info(f'Original {pid=} terminating {p}')
            if IS_WINDOWS:
                # I would like to call:
                #     os.kill(p.pid, signal.CTRL_BREAK_EVENT)
                # But that only works if I have a tty and most Windows python to
                # check isatty() is often unrealiable. We will instead skip to
                # taskkill later
                continue

            # else, non-windows:
            p.terminate() # SIGTERM

        except (ProcessLookupError, psutil.NoSuchProcess):
            pass

    # 4. Wait and ungraceful killing (kill / SIGKILL):
    _, alive = psutil.wait_procs(all_to_kill.values(), timeout=3)
    if alive:
        if IS_WINDOWS:
            # Windows, using subprocess to do the killing:
            # /F = Force, /T = Tree
            subprocess.run( # pylint: disable=subprocess-run-check
                ['taskkill', '/F', '/T', '/PID', str(pid)],
                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
            )
        else:
            # non-Windows, target the PGID group to kill escaped children
            if pgid:
                try:
                    info(f'Original {pid=} killing {pgid=}')
                    # pylint: disable=no-member
                    os.killpg(pgid, signal.SIGKILL)
                except Exception:
                    pass

            # non-Windows, final pass on any PIDs psutil still sees:
            for p in alive:
                if not p:
                    continue
                try:
                    info(f'Original {pid=} killing survivor {p}')
                    p.kill()
                except (ProcessLookupError, psutil.NoSuchProcess):
                    pass
