"""This module defines a function to run a command in a temporary directory,
with specified input files present, and captures the contents of specified
output files.

This is the simplest building block used as the base for rdfox_runner.

"""


from io import TextIOBase
import os
import logging
import subprocess
import shutil
from pathlib import Path
from tempfile import mkdtemp
import threading
import time

from typing import Any, Callable, List, Union, Optional, TextIO, Mapping

StrPath = Union[str, Path]
PathOrIO = Union[StrPath, TextIO]

logger = logging.getLogger(__name__)


class CommandRunner:
    """Run a command in a temporary directory.

    This can be used as a context manager, ensuring the temporary directory is
    cleaned up and the subprocess is stopped when finished with.

    :param input_files: mapping with keys being the target path and value the
        source path.
    :param command: command to run, as passed to :class:`subprocess.Popen`
    :param shell: whether to run command within shell
    :param wait_before_enter: whether to wait for command to complete before
        continuing with context manager body.
    :param wait_before_exit: whether to wait for command to complete by itself
        before terminating it, when leaving context manager body.
    :param timeout: timeout if `wait` is true.
    :param working_dir: Path to setup command in, defaults to a temporary
        directory
    :param output_callback: Callback on output from command.
    """

    def __init__(
            self,
            input_files: Optional[Mapping[str, PathOrIO]] = None,
            command: Optional[Union[List, str]] = None,
            shell: bool = False,
            wait_before_enter: bool = False,
            wait_before_exit: bool = False,
            timeout: Optional[float] = None,
            working_dir: Optional[StrPath] = None,
            output_callback: Optional[Callable] = None,
        ):

        if working_dir is not None:
            working_dir = Path(working_dir)
            if working_dir.exists() and any(working_dir.iterdir()):
                logger.warning(f"Existing working directory not empty: {working_dir}")

        self.input_files = input_files or {}
        self.command = command
        self.shell = shell
        self.wait_before_enter = wait_before_enter
        self.wait_before_exit = wait_before_exit
        self.timeout = timeout
        self.working_dir = working_dir
        self.output_callback = output_callback

        self._process = None
        self._cleanup_working_dir = (working_dir is None)

        if os.environ.get("RDFOX_RUNNER_KEEP_WORKING_DIR", ""):
            self._cleanup_working_dir = False

    def start(self):
        """Setup files and start the command running.

        This is a convenience method to run :meth:`setup_files` and
        :meth:`start_subprocess` together, as needed.
        """
        self.setup_files()
        if self.command:
            self.start_subprocess()

    def stop(self):
        """Stop the command and clean up files.

        This is a convenience method to run :meth:`stop_subprocess` and
        :meth:`cleanup_files` together, as needed.

        :raises subprocess.CalledProcessError: if the subprocess returns an
            error exit code.
        """
        if self.command:
            self.stop_subprocess()
        if self._cleanup_working_dir:
            self.cleanup_files()

    def setup_files(self):
        """Setup the files ready to run the command.

        If :attr:`working_dir` has been specified, it is created if it does not
        exist.  Otherwise, a new temporary directory is created.

        The files listed in :attr:`input_files` are copied into the working
        directory.
        """
        if self.working_dir is None:
            self.working_dir = Path(mkdtemp())
            logger.debug("Created temporary working directory %s", self.working_dir)
        elif not self.working_dir.exists():
            self.working_dir.mkdir(parents=True)
            logger.debug("Created temporary working directory %s", self.working_dir)

        logger.info("Setting up to run command in %s", self.working_dir)

        for target, source in self.input_files.items():
            copy_files(source, self.working_dir / target)

    def start_subprocess(self):
        """Start the subprocess running.

        If :attr:`wait` is true, wait for up to :attr:`timeout` seconds before
        continuing.
        """
        if self.command is None:
            return

        logger.info("Running command '%s'", self.command)

        self._process = subprocess.Popen(
            self.command,
            cwd=self.working_dir,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            # text=True,
            # bufsize=1,  # line buffered
            shell=self.shell,
        )

        self._output_thread = threading.Thread(target=output_reader,
                                               args=(self._process, self.output_callback))
        self._output_thread.daemon = True
        self._output_thread.start()

        logger.debug("finished starting")

    def wait(self):
        """Wait for subprocess to exit.

        Waits up to :attr:`timeout` seconds.
        """
        if self.timeout is not None:
            logger.debug("waiting %.1f s for subprocess to finish...", self.timeout)
        else:
            logger.debug("waiting for ever for subprocess to finish...")
        try:
            self._process.wait(timeout=self.timeout)
            logger.info('Subprocess exited with returncode = %s', self._process.returncode)
        except subprocess.TimeoutExpired:
            logger.error('Subprocess did not terminate in time when waiting, continuing')
        except KeyboardInterrupt:
            logger.info('KeyboardInterrupt: stopping subprocess...')
            self.stop_subprocess()
            raise

    def stop_subprocess(self):
        """Stop the subprocess. """
        if self.command is None:
            return

        if not self._process:
            logger.warning("trying to stop process that wasn't started")
            return

        logger.debug("trying to terminate processs...")
        try:
            self._process.stdin.close()
        except (OSError, BrokenPipeError):
            # On Windows it's an OSError, see https://bugs.python.org/issue35754
            pass
        self._process.terminate()
        logger.debug("...terminate returned.")
        try:
            logger.debug("waiting for process to exit...")
            self._process.wait(timeout=5)
            logger.info('Subprocess exited with returncode = %s', self._process.returncode)
        except subprocess.TimeoutExpired:
            logger.error('Subprocess did not terminate in time')
            self._process.kill()

        self._output_thread.join()

        self._process.poll()  # update returncode
        if self._process.returncode > 0:
            logger.error("Error running command: %d", self._process.returncode)
            # raise subprocess.CalledProcessError(self._process.returncode, cmd=self.command)
        elif self._process.returncode < 0:
            logger.warning("Process was killed: returncode=%d", self._process.returncode)

    @property
    def returncode(self):
        if self._process:
            return self._process.returncode
        return None

    def cleanup_files(self):
        """Cleanup temporary working directory, if needed.

        The directory is only removed if it was newly created, not if it was
        passed in as :attr:`working_dir`.
        """
        if self.working_dir is not None:
            if not self._cleanup_working_dir:
                logger.warning("trying to cleanup working directory that wasn't created")
                return

            shutil.rmtree(self.working_dir, ignore_errors=True)
            logger.debug(f"Removed temporary working directory {self.working_dir}")
            self.working_dir = None

    def __enter__(self):
        self.start()
        if self.wait_before_enter:
            self.wait()
        return self

    def __exit__(self, exc, value, tb):
        if self.wait_before_exit:
            self.wait()
        self.stop()

    def files(self, path) -> Path:
        """Return path to temporary directory.

        :param path: path relative to the working directory
        """
        return self.working_dir / path


def copy_files(src: PathOrIO, dst: Path):
    if isinstance(src, (Path, str)):
        src = Path(src)
        if src.is_dir():
            if dst.exists():
                shutil.rmtree(dst)
            logger.debug("Copying directory %s to %s", src, dst)
            start_time = time.perf_counter()
            shutil.copytree(src, dst)
            logger.debug("Finished in %d ms", (time.perf_counter() - start_time) * 1000)
        else:
            dst.parent.mkdir(parents=True, exist_ok=True)
            logger.debug("Copying file %s to %s", src, dst)
            start_time = time.perf_counter()
            shutil.copy(src, dst)
            logger.debug("Finished in %d ms", (time.perf_counter() - start_time) * 1000)
    else:
        dst.parent.mkdir(parents=True, exist_ok=True)
        logger.debug("Writing data %s", dst)
        start_time = time.perf_counter()
        mode = "t" if isinstance(src, TextIOBase) else "b"
        with open(dst, "w" + mode) as f:
            f.write(src.read())
        logger.debug("Finished in %d ms", (time.perf_counter() - start_time) * 1000)


def get_file_contents(root_path: Path, output_files: Mapping[Any, StrPath]):
    """Read the contents of output_files.

    :param root_path: Root path that output_files are relative to
    :param output_files: Dict of {label: target path} to collect
    :return: Dict of {label: result}
    """
    result = {}
    for label, rel_path in output_files.items():
        with open(root_path / rel_path, "rt") as f:
            result[label] = f.read()
    return result


def output_reader(process, output_callback):
    logger.debug("output reader started (output_callback=%s)", output_callback)
    for line in iter(process.stdout.readline, b""):
        line = line.decode("utf-8").rstrip()
        # line = line.rstrip()
        logger.debug("cmd> %s", line)
        if output_callback:
            output_callback(line)
