# Copyright (C) 2018  Lars Wirzenius
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.


import logging
import os
import subprocess


import cliapp


class Runner:

    def __init__(self, reporter, maxbuf=128*1024):
        self._reporter = reporter
        self._buffers = {
            'stdout': b'',
            'stderr': b'',
        }
        self._maxbuf = maxbuf
        self._timeout = 1.0

    def runcmd(self, *argvs, **kwargs):
        for argv in argvs:
            logging.debug('Runner.runcmd: argv: %r', argv)
        for key in kwargs:
            logging.debug('Runner.runcmd: kwargs: %r=%r', key, kwargs[key])
        assert all(argv is not None for argv in argvs)
        with open('/dev/null', 'rb') as devnull:
            exit_code, _, _ = cliapp.runcmd_unchecked(
                *argvs,
                stdin=devnull,
                stdout_callback=self.capture_stdout,
                stderr=subprocess.STDOUT,
                output_timeout=self._timeout,
                timeout_callback=self.flush,
                **kwargs
            )
        self.flush()
        logging.debug('Runner.runcmd: finished, exit code: %d', exit_code)
        return exit_code

    def capture_stdout(self, data):
        return self.capture('stdout', data)

    def capture(self, stream_name, data):
        self._buffers[stream_name] += data
        if len(self._buffers[stream_name]) >= self._maxbuf:
            self.flush()

        return b''

    def flush(self):
        stdout = self._buffers['stdout']
        stderr = self._buffers['stderr']
        self._reporter.report(None, stdout, stderr)
        self._buffers['stdout'] = b''
        self._buffers['stderr'] = b''


class Mounter:  # pragma: no cover

    def __init__(self, mounts, runner):
        self._mounts = mounts
        self._runner = runner

    def __enter__(self):
        self.mount()
        return self

    def __exit__(self, *args):
        self.unmount()

    def mount(self):
        for dirname, mp in self._mounts:
            if not os.path.exists(mp):
                os.mkdir(mp)
            self._runner.runcmd(['sudo', 'mount', '--bind', dirname, mp])

    def unmount(self):
        for dirname, mp in reversed(self._mounts):
            try:
                self._runner.runcmd(['sudo', 'umount', mp])
            except BaseException as e:
                logging.error(
                    'Ignoring error while unmounting %s: %s', mp, str(e))


class ActionEnvironment:  # pragma: no cover

    def __init__(self, systree, workspace, reporter):
        super().__init__()
        self._systree = systree
        self._workspace = workspace
        self._reporter = reporter
        self._extra_env = {}

    def set_extra_env(self, extra_env):
        self._extra_env = dict(extra_env)

    def get_systree_directory(self):
        return self._systree

    def get_workspace_directory(self):
        return self._workspace

    def get_mounts(self):
        return []

    def report(self, exit_code, msg):
        self._reporter.report(exit_code, msg, None)

    def runcmd(self, argv):
        raise NotImplementedError()

    def host_runcmd(self, *argvs, cwd=None):
        env = self.get_env_vars()
        runner = Runner(self._reporter)
        mounts = self.get_mounts()
        with Mounter(mounts, runner):
            return runner.runcmd(*argvs, cwd=cwd, env=env)

    def get_env_vars(self):
        return dict(self._extra_env)


class HostEnvironment(ActionEnvironment):

    def runcmd(self, argv):
        return self.host_runcmd(argv, cwd=self._workspace)


class ChrootEnvironment(ActionEnvironment):

    def runcmd(self, argv):
        prefix = ['sudo', 'chroot', self._workspace]
        return self.host_runcmd(prefix + argv)


class ContainerEnvironment(ActionEnvironment):

    def runcmd(self, argv):
        bind = '{}:/workspace'.format(self._workspace)
        prefix = [
            'sudo', 'systemd-nspawn',
            '-D', self._systree,
            '--bind', bind,
            '--chdir', '/workspace',
        ]
        env = self.get_env_vars()
        for key in env:  # pragma: no cover
            var = '{}={}'.format(key, env[key])
            prefix.extend(['--setenv', var])
        return self.host_runcmd(prefix + argv)