#!/usr/bin/python3
#
# Copyright (c) 2018  Peter Pentchev <roam@ringlet.net>
#
# This file conforms to the GNU General Public License version 2.

import os
import pty
import re
import subprocess
import unittest

import ddt


class ProcessDriver(object):
    """A base class for drivers for running subprocesses."""

    def __init__(self, testCase):
        """Initialize a driver object.  Stores the ```testCase```."""
        self.testCase = testCase

    def start_process(self, cmd, argv):
        """Start a child process and run the specified command."""
        raise Exception('{cls}.start_process() must be overridden'
                        .format(cls=type(self).__name__))

    def read_data(self):
        """Read some data from the child process's stdout."""
        raise Exception('{cls}.read_data() must be overridden'
                        .format(cls=type(self).__name__))

    def write_data(self, data):
        """Write some data from the child process's stdin."""
        raise Exception('{cls}.write_data() must be overridden'
                        .format(cls=type(self).__name__))

    def check_exit_code(self, expected):
        """Check whether the command exited as expected."""
        raise Exception('{cls}.check_exit_code() must be overridden'
                        .format(cls=type(self).__name__))


class PtyDriver(ProcessDriver):
    """Run commands attached to a pseudo-tty."""

    def __init__(self, testCase):
        """Clear out the subprocess pid and fd."""
        super(PtyDriver, self).__init__(testCase)
        self.pid = None
        self.fd = None

    def start_process(self, cmd, argv):
        """Start a child process attached to a pseudo-tty."""
        self.testCase.assertIsNone(self.pid)
        self.testCase.assertIsNone(self.fd)

        (pid, fd) = pty.fork()
        self.testCase.assertIsNotNone(pid)
        if pid == 0:
            os.execlp(cmd, *argv)
            os._exit(42)

        self.testCase.assertIsNotNone(fd)
        self.pid = pid
        self.fd = fd

    def read_data(self):
        """Read some data from the pseudo-tty."""
        self.testCase.assertIsNotNone(self.fd)
        return os.read(self.fd, 1024)

    def write_data(self, data):
        """Write some data to the pseudo-tty."""
        os.write(self.fd, bytes(data, 'US-ASCII'))

    def check_exit_status(self, expected):
        """Wait for the child process, check the exit code."""
        (npid, status) = os.waitpid(self.pid, 0)
        self.testCase.assertEqual(self.pid, npid)
        self.testCase.assertEqual(status, 256 * expected)


class TmuxDriver(ProcessDriver):
    """Run a process in a new tmux session."""
    def __init__(self, testCase):
        """Initialize a tmux driver object."""
        super(TmuxDriver, self).__init__(testCase)
        self.p = None
        self.output = None
        self.in_decoded = ''
        self.re_output = re.compile('%output \s+ \S+ \s+ (?P<data> .*)', re.X)
        self.re_output_decode = re.compile(
            '(?P<pre> .*? ) [\\\\] (?P<code> [0-9]{3} ) (?P<post> .* )',
            re.X)

    def read_line(self):
        """Read and decode a single line from tmux."""

        while self.in_decoded.find('\n') == -1:
            raw = self.p.stdout.read1(1024)
            if len(raw) == 0:
                if self.p.poll() is not None:
                    line = self.in_decoded
                    self.in_decoded = ''
                    return line
            else:
                self.in_decoded += raw.decode('US-ASCII')

        idx = self.in_decoded.find('\n')
        assert(idx != -1)
        line = self.in_decoded[:idx]
        self.in_decoded = self.in_decoded[idx + 1:]
        if len(line) > 0 and line[-1] == '\r':
            line = line[:-1]
        return line

    def start_process(self, cmd, argv):
        """Start a tmux process."""
        self.testCase.assertIsNone(self.p)
        p = subprocess.Popen(['tmux', '-C', 'new-session', '--'] + argv,
                             stdin=subprocess.PIPE,
                             stdout=subprocess.PIPE,
                             shell=False)
        self.p = p

        # Wait for the window
        while True:
            line = self.read_line()
            self.testCase.assertNotEqual(line, '')
            words = line.split(' ')
            self.testCase.assertNotEqual(words, [])
            self.testCase.assertRegex(words[0], '%')
            self.testCase.assertNotEqual(words[0], '%output')
            if words[0] == '%window-add':
                self.testCase.assertEqual(len(words), 2)
                self.window = words[1]
                break

    def check_output(self, line):
        """Check if a line is a %output notification."""
        m = self.re_output.match(line)
        if m is None:
            return None
        return self.decode_output(m.groupdict()['data'])

    def decode_output(self, output):
        res = ''
        while True:
            m = self.re_output_decode.match(output)
            if m is None:
                return res + output
            d = m.groupdict()
            res += d['pre'] + bytes([int(d['code'], 8)]).decode('US-ASCII')
            output = d['post']

    def read_data(self):
        """Return data received from %output lines."""
        if self.output is not None:
            res = self.output
            self.output = None
            return bytes(res, 'US-ASCII')

        while True:
            line = self.read_line()
            if line == '':
                raise OSError()
            output = self.check_output(line)
            if output is not None:
                return bytes(output, 'US-ASCII')

    def send_keys(self, data):
        cmd = "send-keys -t {win} {data}\n".format(win=self.window, data=data)
        self.p.stdin.write(bytes(cmd, 'US-ASCII'))
        self.p.stdin.flush()

    def write_quoted(self, word):
        self.send_keys("-l '{word}'".format(word=word))

    def write_single_quote(self):
        self.send_keys('-l "' + "'" + '"')

    def write_enter(self):
        self.send_keys('enter')

    def write_single_line(self, line):
        try:
            while True:
                idx = line.index('\n')
                word = line[:idx]
                line = line[idx + 1:]
                self.write_quoted(word)
                self.write_single_quote()
        except ValueError:
            self.write_quoted(line)

    def write_data(self, data):
        """Encode and write out the data."""
        try:
            while True:
                idx = data.index('\n')
                line = data[:idx]
                data = data[idx + 1:]
                self.write_single_line(line)
                self.write_enter()
        except ValueError:
            if len(data) > 0:
                self.write_single_line(data)

    def check_exit_status(self, expected):
        """Just fake it..."""
        self.p.stdin.close()
        self.p.wait()
        self.p.stdout.close()


def read_prompt(driver, check=True):
    """
    Read a series of characters ending in ": " and return it.
    If `check` is false, do not look for ": ", but read to
    the end of the stream and return everything.
    """
    buf = ''
    while True:
        try:
            s = driver.read_data().decode()
        except OSError:
            return buf
        buf = buf + s
        if check and buf.endswith(': '):
            return buf


@ddt.ddt
class TestDonkey(unittest.TestCase):
    drivers = {
        'pty': PtyDriver,
        'tmux': TmuxDriver,
    }

    @ddt.data(
        {
            'argv': ['-h'],
            'dialog': [
            ],
            'result': 'usage:\tdonkey',
        },

        {
            'argv': ['--help'],
            'dialog': [
            ],
            'result': 'usage:\tdonkey',
        },

        {
            'argv': ['-v'],
            'dialog': [
            ],
            'result': 'donkey version ',
        },

        {
            'argv': ['--version'],
            'dialog': [
            ],
            'result': 'donkey version ',
        },

        {
            'argv': ['--features'],
            'dialog': [
            ],
            'result': 'Features: donkey=',
        },

        {
            'argv': ['-X'],
            'dialog': [
            ],
            'result': 'usage:\tdonkey',
            'exit_code': 1,
        },

        {
            'argv': ['-i'],
            'dialog': [
                ('Enter login name', 'roam'),
                ('Enter sequence', '99'),
                ('Enter new seed', 'st66708'),
                ('Enter passphrase', 'meowmeow'),
                ('Re-enter passphrase', 'meowmeow'),
            ],
            'result': 'FOLD BAWD AS WARM FOUL BARN',
        },

        {
            'argv': ['24', 'st66708'],
            'dialog': [
                ('Enter passphrase', 'meowmeow'),
            ],
            'result': 'INCH CLAD LOCK VIE PAL WINK',
        },

        {
            'argv': ['-n', '10', '77', 'st66708'],
            'dialog': [
                ('Enter passphrase', 'meowmeow'),
            ],
            'result': '73: HYMN GRAD THEY MANY MUTT BONG',
        },

        {
            'argv': ['-i', '-f', 'md4'],
            'dialog': [
                ('Enter login name', 'roam'),
                ('Enter sequence', '99'),
                ('Enter new seed', 'st66708'),
                ('Enter passphrase', 'meowmeow'),
                ('Re-enter passphrase', 'meowmeow'),
            ],
            'result': 'FOLD BAWD AS WARM FOUL BARN',
        },

        {
            'argv': ['-f', 'md4', '24', 'st66708'],
            'dialog': [
                ('Enter passphrase', 'meowmeow'),
            ],
            'result': 'INCH CLAD LOCK VIE PAL WINK',
        },

        {
            'argv': ['-n', '10', '-f', 'md4', '77', 'st66708'],
            'dialog': [
                ('Enter passphrase', 'meowmeow'),
            ],
            'result': '73: HYMN GRAD THEY MANY MUTT BONG',
        },

        {
            'argv': ['-i', '-f', 'md5'],
            'dialog': [
                ('Enter login name', 'roam'),
                ('Enter sequence', '99'),
                ('Enter new seed', 'st66708'),
                ('Enter passphrase', 'meowmeow'),
                ('Re-enter passphrase', 'meowmeow'),
            ],
            'result': 'JILT LURA LULU HOME WERE SWAM',
        },

        {
            'argv': ['-f', 'md5', '24', 'st66708'],
            'dialog': [
                ('Enter passphrase', 'meowmeow'),
            ],
            'result': 'SEND OR LOOT BETH TONY WERT',
        },

        {
            'argv': ['-n', '10', '-f', 'md5', '77', 'st66708'],
            'dialog': [
                ('Enter passphrase', 'meowmeow'),
            ],
            'result': '73: BUDD DOLE LIT BUNT BLOW ADD',
        },
    )
    def test_dialog(self, data):
        """
        Start the test program with the options given in the test data,
        carry out the dialog, then check the result.
        """
        driver_name = os.environ.get('TEST_PROCESS_DRIVER', 'tmux')
        self.assertIn(driver_name, self.drivers)
        driver = self.drivers[driver_name](testCase=self)
        cmd = os.environ.get('TEST_PROG', './donkey')
        driver.start_process(cmd, [cmd] + data['argv'])

        for (prompt, response) in data['dialog']:
            p = read_prompt(driver)
            self.assertIn(prompt, p)
            driver.write_data(response + '\n')

        p = read_prompt(driver, False)
        self.assertIn(data['result'], p)

        driver.check_exit_status(data.get('exit_code', 0))


if __name__ == '__main__':
    unittest.main()
