IoTVirtualAssistant/utils/prompt.py

import dotenv
import getpass
import platform
import os
from prompt_toolkit.shortcuts import prompt
from prompt_toolkit import print_formatted_text, HTML
from prompt_toolkit.styles import Style
from prompt_toolkit.cursor_shapes import CursorShape
from prompt_toolkit.shortcuts import button_dialog
from prompt_toolkit.application import run_in_terminal
from prompt_toolkit.key_binding import KeyBindings
from prompt_toolkit import PromptSession
from prompt_toolkit.history import FileHistory


dotenv.load_dotenv()

# bindings = KeyBindings()

# @bindings.add('c-m')
# def _(event):
#     " Say 'hello' when `c-t` is pressed. "
#     def print_hello():
#         print('hello world')
#     run_in_terminal(print_hello)

# @bindings.add('c-x')
# def _(event):
#     " Exit when `c-x` is pressed. "
#     event.app.exit()


class PromptTheme:
    def __init__(self, appname):
        self.message = [
            ('class:username', getpass.getuser()),
            ('class:at',       '@'),
            ('class:host',     platform.node()),
            ('class:colon',    ':'),
            ('class:appname',    appname),
            ('class:pound',    '# '),
        ]
        self.style = Style.from_dict({
                    '':          '#ffffff',
                    'username': '#884444',
                    'at':       '#00aa00',
                    'colon':    '#ffffff',
                    'pound':    '#00aa00',
                    'appname':     '#00ffff bg:#444400',
                    'path':     'ansicyan underline',
                    'info':    'black bg:white',
                    'warning':    'yellow bg:white',
                    'error':    'red bg:white',
                    'bottom-toolbar': '#ffffff bg:#333333',
                    
                })

        if os.getenv('PROMPT_THEME'):
            theme = os.getenv('PROMPT_THEME')
            if theme == 'dark_plus':
                self.style = Style.from_dict({
                    '':          '#ffffff',
                    'username': '#884444',
                    'at':       '#00aa00',
                    'colon':    '#ffffff',
                    'pound':    '#00aa00',
                    'appname':     '#00ffff bg:#444400',
                    'path':     'ansicyan underline',
                    'info':    'black bg:white',
                    'warning':    'red bg:yellow',
                    'error':    'red bg:white',
                    'bottom-toolbar': '#ffffff bg:#333333',
                })



class Prompt:
    def __init__(self, appname='chitchat'):
        theme = PromptTheme(appname)
        self.message = theme.message
        self.style = theme.style
        self.bottom_toolbar_text = ' Press [Ctrl+X] to exit '

        # if history file is not present, create it
        history_file = os.path.expanduser('~/.aiassistant_history')
        if not os.path.exists(history_file):
            with open(history_file, 'w') as f:
                pass

                
        self.session = PromptSession(history=FileHistory(history_file))

    def bottom_toolbar(self):
        return [('class:bottom-toolbar', f"@aiassistant:~ {self.bottom_toolbar_text}")]


    def print(self, message, tag='info'):
        print_formatted_text(HTML(f"<{tag}>{message}</{tag}>"), style=self.style)


    def get_input_mode(self):
        return button_dialog(
            title='Input Mode',
            text='What input mode do you want to use?',
            buttons=[
                ('Text', 'text'),
                ('Voice', 'voice'),
                ('Surprise me', None)
            ],
        ).run()
    
    def get(self):
        return self.session.prompt(
            self.message, 
            style=self.style, 
            cursor=CursorShape.BLINKING_BLOCK, 
            bottom_toolbar=self.bottom_toolbar,
            )