#!/usr/bin/env python3

import os
import sys
import time
import json
import shlex
import yaml
import queue
import argparse
import threading
import subprocess
from collections import OrderedDict

from shells import BashShell, FishShell, ZshShell

# =============================================================================
# Script configuration
# =============================================================================

SHELLS              = ['bash', 'fish', 'zsh']
TMUX_SESSION_PREFIX = 'crazy-complete-test'
TESTS_INFILE        = 'tests.yaml'
TESTS_OUTFILE       = 'tests.new.yaml'
CRAZY_COMPLETE      = '../../crazy-complete'
COMPLETIONS_OUTDIR  = 'output'
TEMP_DIR            = 'tmp'
OUTPUT_WAIT         = 5

# =============================================================================
# Switch to the script's directory
# =============================================================================

os.chdir(os.path.dirname(os.path.abspath(__file__)))
CURRENT_DIRECTORY = os.path.abspath(os.getcwd())

# =============================================================================
# Commandline parser
# =============================================================================

argp = argparse.ArgumentParser()
argp.add_argument('-f', '--fast', action='store_true', default=False,
    help='Enable fast testing mode. For tests where the input matches the expected output, these tests will always pass')
argp.add_argument('-d', '--driver', default='tmux', choices=['pyte', 'tmux'],
    help='Select driver for tests')
argp.add_argument('-t', '--threads', default=1, type=int,
    help='Set the number of threads per shell')
argp.add_argument('-w', '--wait', default=OUTPUT_WAIT, type=float,
    help='Set how many seconds to wait for terminal output')
opts = argp.parse_args()

OUTPUT_WAIT = opts.wait

# =============================================================================
# Helper functions
# =============================================================================

def print_err(*args):
    print(*args, file=sys.stderr)

def run(args, env=None):
    result = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, env=env)

    if result.returncode != 0:
        raise Exception("Command %r failed: %s" % (args, result.stderr))

    return result.stdout

def indent(string, num_spaces):
    lines = string.split('\n')
    indented_lines = [((' ' * num_spaces) + line) if line.strip() else line for line in lines]
    return '\n'.join(indented_lines)

def make_yaml_block_string(s):
    return "|\n%s" % indent(s, 2)

# =============================================================================
# Import driver depending on command line arguments
# =============================================================================

if opts.driver == 'tmux':
    from tmux_driver import TmuxTerminal
elif opts.driver == 'pyte':
    print_err("The `pyte` driver has been disabled because it does not implement the `kitty`")
    print_err("protocol used by the Fish shell, causing tests to fail.")
    sys.exit(2)

# =============================================================================
# Ensure all dependencies are available
# =============================================================================

for shell in SHELLS:
    try:
        run(['sh', '-c', f'type {shell}'])
    except Exception:
        print_err(f'Program `{shell}` not found')
        sys.exit(2)

if not os.path.exists('/usr/share/bash-completion/bash_completion'):
    print_err('File `/usr/share/bash-completion/bash_completion` not found. Is `bash-completion` installed?')
    sys.exit(2)

# =============================================================================
# Test code
# =============================================================================

class Test:
    def __init__(self, data):
        self.number = data.get('number', None) 
        self.definition_file = data['definition_file']
        self.description = data['description']
        self.comment = data.get('comment', None)
        self.send = data['send']
        self.require_order = data.get('require_order', False)
        self.bash_tabs = data.get('bash_tabs', 1)
        self.fish_tabs = data.get('fish_tabs', 1)
        self.zsh_tabs = data.get('zsh_tabs', 1)
        self.bash_expected = data.get('bash_expected', '').strip()
        self.fish_expected = data.get('fish_expected', '').strip()
        self.zsh_expected = data.get('zsh_expected', '').strip()
        self.bash_result = None
        self.fish_result = None
        self.zsh_result = None

    def get_tabs(self, shell):
        if shell == 'bash':
            return self.bash_tabs
        if shell == 'fish':
            return self.fish_tabs
        if shell == 'zsh':
            return self.zsh_tabs
        raise

    def get_expected(self, shell):
        if shell == 'bash':
            return self.bash_expected
        if shell == 'fish':
            return self.fish_expected
        if shell == 'zsh':
            return self.zsh_expected

    def get_result(self, shell):
        if shell == 'bash':
            return self.bash_result
        if shell == 'fish':
            return self.fish_result
        if shell == 'zsh':
            return self.zsh_result
        raise

    def set_result(self, shell, result):
        if shell == 'bash':
            self.bash_result = result
        elif shell == 'fish':
            self.fish_result = result
        elif shell == 'zsh':
            self.zsh_result = result
        else:
            raise

    def passed(self, shell):
        result = self.get_result(shell)
        expected = self.get_expected(shell)

        if self.require_order:
            return result.split() == expected.split()
        else:
            return sorted(result.split()) == sorted(expected.split())

    def to_yaml(self):
        r = OrderedDict()

        r['number']          = str(self.number)
        r['definition_file'] = json.dumps(self.definition_file)
        r['description']     = json.dumps(self.description)
        if self.comment:
            r['comment']     = json.dumps(self.comment)
        r['send']            = json.dumps(self.send)
        if self.require_order:
            r['require_order'] = 'true'
        if self.bash_tabs != 1:
            r['bash_tabs']   = str(self.bash_tabs)
        r['bash_expected']   = make_yaml_block_string(self.bash_result)
        if self.fish_tabs != 1:
            r['fish_tabs']   = str(self.fish_tabs)
        r['fish_expected']   = make_yaml_block_string(self.fish_result)
        if self.zsh_tabs != 1:
            r['zsh_tabs']    = str(self.zsh_tabs)
        r['zsh_expected']    = make_yaml_block_string(self.zsh_result)

        return '\n'.join('%s: %s' % key_value for key_value in r.items())


class Tests:
    def __init__(self, tests):
        self.definition_files = tests.pop(0)
        self.tests = [Test(data) for data in tests]

    def enumerate_tests(self):
        number = 1
        for test in self.tests:
            test.number = number
            number += 1

    def find_test_by_number(self, num):
        for test in self.tests:
            if test.number == num:
                return test

        raise Exception(f"No test with number {num} found")

    def write_tests_file(self, outfile):
        r = []

        files = ''
        for file, args in self.definition_files.items():
            files += '%s: %s\n' % (file, json.dumps(args))
        r += [files.strip()]

        for test in self.tests:
            r += [test.to_yaml()]

        r = '\n---\n'.join(r)

        with open(outfile, 'w') as fh:
            fh.write(r)

    def generate_completion_files(self):
        try:
            os.mkdir(COMPLETIONS_OUTDIR)
        except FileExistsError:
            if not os.path.isdir(COMPLETIONS_OUTDIR):
                raise NotADirectoryError(COMPLETIONS_OUTDIR) from None

        try:
            os.mkdir(TEMP_DIR)
        except FileExistsError:
            if not os.path.isdir(TEMP_DIR):
                raise NotADirectoryError(TEMP_DIR) from None

        batch_file = '%s/%s' % (TEMP_DIR, 'batch.txt')
        batch_commands = []

        for file, args in self.definition_files.items():
            definition_file = args['definition_file']
            temp_definition_file = f'{TEMP_DIR}/{definition_file}'

            with open(definition_file, 'r', encoding='UTF-8') as fh:
                content = fh.read()
                content = content.replace('%TEST_DIR%', CURRENT_DIRECTORY)

            with open(temp_definition_file, 'w', encoding='UTF-8') as fh:
                fh.write(content)

            for shell in SHELLS:
                cmd =  ['--debug', '--zsh-compdef=False']
                for arg in args['args']:
                    cmd += [arg.replace('$shell', shell)]
                cmd += ['-o', f'{COMPLETIONS_OUTDIR}/{file}.{shell}']
                cmd += [shell, temp_definition_file]
                cmd = ' '.join(shlex.quote(arg) for arg in cmd)
                batch_commands.append(cmd)

        with open(batch_file, 'w', encoding='utf-8') as fh:
            print_err("Writing batch file '%s' ..." % batch_file)
            fh.write('\n'.join(batch_commands))

        print_err('Running batch file ...')
        args = [CRAZY_COMPLETE, '--batch', batch_file]
        subprocess.run(args, check=True)
        print_err('Finished running batch file ...')

def tests_worker_thread(thread_id, shell, input_queue, result_queue):
    if opts.driver == 'tmux':
        term = TmuxTerminal(f'{TMUX_SESSION_PREFIX}-{shell}-{thread_id}')
    elif opts.driver == 'pyte':
        term = PyteTerminal()

    term_shell = {
        'bash': BashShell,
        'fish': FishShell,
        'zsh':  ZshShell
    }[shell](term)

    old_definition_file = None

    while True:
        try:
            test = input_queue.get_nowait()
        except queue.Empty:
            break

        if old_definition_file != test.definition_file:
            old_definition_file = test.definition_file

            try:
                term_shell.stop()
            except Exception:
                pass
            completion_file = '%s/%s.%s' % (COMPLETIONS_OUTDIR, test.definition_file, shell)
            term_shell.start()
            term.resize_window(80, 100)
            term_shell.set_prompt()
            term_shell.init_completion()
            term_shell.load_completion(completion_file)
            time.sleep(1.5)

        output = term.complete(
            test.send,
            test.get_tabs(shell),
            OUTPUT_WAIT,
            test.get_expected(shell),
            opts.fast)

        result = {
            'number': test.number,
            'shell':  shell,
            'result': output
        }

        result_queue.put(result)

    try:
        term_shell.stop()
    except Exception:
        pass

class Tester():
    def __init__(self, tests):
        self.tests = tests
        self.result_queue = queue.Queue()
        self.input_queues = {}
        self.threads = []
        self.num_failed = 0

    def run(self):
        for shell in SHELLS:
            self.input_queues[shell] = queue.Queue()
            for test in self.tests.tests:
                self.input_queues[shell].put(test)

        for thread_id in range(opts.threads):
            for shell in SHELLS:
                thread = threading.Thread(
                    target=tests_worker_thread,
                    args=(thread_id, shell, self.input_queues[shell], self.result_queue)
                )

                thread.start()
                self.threads.append(thread)

        while self.threads_are_running():
            self.eat_queue()
            time.sleep(0.1)
        self.eat_queue()

    def threads_are_running(self):
        for thread in self.threads:
            if thread.is_alive():
                return True
        return False

    def eat_queue(self):
        while not self.result_queue.empty():
            self.process_result(self.result_queue.get())

    def process_result(self, result):
        test = self.tests.find_test_by_number(result['number'])
        shell = result['shell']
        test.set_result(shell, result['result'])

        if not test.passed(shell):
            self.num_failed += 1
            print_err("Test #%02d (%-4s - %s) failed" % (test.number, shell, test.description))
        else:
            print_err("Test #%02d (%-4s - %s) OK" % (test.number, shell, test.description))

# =============================================================================
# Main
# =============================================================================

with open(TESTS_INFILE, 'r') as fh:
    tests = Tests(list(yaml.safe_load_all(fh)))

tests.enumerate_tests()
tests.generate_completion_files()
tester = Tester(tests)
tester.run()
tests.write_tests_file(TESTS_OUTFILE)
if tester.num_failed:
    print_err(f'{tester.num_failed} tests failed.')
    print_err(f'Use diff or vimdiff on `{TESTS_INFILE}` and `{TESTS_OUTFILE}` for further details')
    if opts.threads > 1:
        print_err('NOTE:')
        print_err(' A high value for -t|--threads may cause that tests fail.')
        print_err(' Consider running this script again with `-t 1`.')
    sys.exit(1)
else:
    print_err("All tests passed")
    sys.exit(0)
