#!/usr/bin/python3

# Copyright (C) 2014-2021 Daniel Baumann <daniel.baumann@open-infrastructure.net>
#
# SPDX-License-Identifier: GPL-3.0+
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from prompt_toolkit import PromptSession

from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
from prompt_toolkit.completion import NestedCompleter
from prompt_toolkit.history import FileHistory
from prompt_toolkit.key_binding import KeyBindings
from prompt_toolkit.shortcuts import set_title
from prompt_toolkit.styles import Style

from pathlib import Path
from subprocess import Popen, PIPE

import os
import socket

home = str(Path.home())
hostname = socket.getfqdn()

container_completer = NestedCompleter.from_nested_dict({
    'auto':     { '--force': None, '--start': None, '--stop': None, },
    'console':  { '--name': None, },
    'create':   { '--name': None, '--capability': None, '--drop-capability': None, '--script': None, '--verbose': None, '--bind': None, '--bind-ro': None, },
    'enter':    { '--name': None, },
    'key':      { '--add': None, '--list': None, '--remove': None, },
    'limit':    { '--name': None, '--blockio-device-weight': None, '--blockio-read-bandwidth': None, '--blockio-weight': None, '--blockio-write-bandwidth': None, '--cpu-quota': None, '--cpu-shares': None, '--memory-limit': None, '--tasks-max': None, },
    'list':     { '--all': None,
                  '--format': {'cli': None, 'csv': None, 'json': None, 'nwdiag': None, 'shell': None, 'yaml': None, 'xml': None, },
                  '--host': None, '--other': None, '--started': None, '--stopped': None,
                },
    'log':      { '--date': None, '--name': None, },
    'move':     { '--force': None, '--new':  None, '--old': None, },
    'remove':   { '--force': None, '--name':  None, },
    'restart':  { '--name':  None, },
    'start':    { '--name':  None, '--verbose': None, },
    'status':   { '--name':  None, },
    'stop':     { '--force': None, '--name':  None, },
    'top':      { '--force': None, '--name':  None, },
    'version': None,
})

style = Style.from_dict({
    'completion-menu.completion': 'bg:ansiblack ansiwhite',
    'completion-menu.completion.current': 'bg:ansibrightblack ansiwhite',
    'bottom-toolbar': 'bg:ansiwhite ansibrightblack',
    'scrollbar.background': 'bg:ansibrightblack',
    'scrollbar.button': 'bg:ansiwhite',
})

kb = KeyBindings()

@kb.add('c-space') # Ctrl-Space
def _(event):
    b = event.app.current_buffer

    if b.complete_state:
        b.complete_next()
    else:
        b.start_completion(select_first=False)

def run(commands):
    with Popen(commands, stdout=PIPE, bufsize=1, universal_newlines=True) as p:
        for line in p.stdout:
            print(line, end='')

    if p.returncode != 0:
        raise CalledProcessError(p.returncode, p.args)

def bottom_toolbar():
    return hostname + ' | [Ctrl-Space] for menu | [Ctrl-D] to exit | [Ctrl-R] for history search | help <command> for documentation'

def main():
    set_title('container-shell: ' + hostname)

    print('container-shell: ' + hostname)
    print()

    os.makedirs(home + '/.cache/container', exist_ok=True)
    history = FileHistory(home + '/.cache/container/history')

    session = PromptSession(
                auto_suggest=AutoSuggestFromHistory(),
                bottom_toolbar=bottom_toolbar,
                completer=container_completer,
                complete_while_typing=True,
                history=history,
                enable_history_search=True,
                key_bindings=kb,
                style=style,
            )

    while True:
        try:
            text = session.prompt('container: ')
        except KeyboardInterrupt:
            continue # Ctrl-C
        except EOFError:
            break # Ctrl-D
        else:
            if not text:
                continue

        try:
            command = text.split(None, 1)[0]
        except:
            command = ''

        if command == 'exit' or command == 'logoff' or command == 'logout' or command == 'quit':
            break

        if command != 'help' and not os.path.exists('/usr/lib/open-infrastructure/container/' + command):
            print('\'' + command + '\': - no such command')
            continue

        commands = text.split()

        if command == 'help':
            try:
                manpage = text.split(None, 1)[1]
                manpage = 'container-' + manpage
            except:
                manpage = 'container'

            commands = [ 'man', manpage ]
        else:
            commands.insert(0, '/usr/bin/container')

        print()
        run(commands)
        print()

if __name__ == '__main__':
    main()