# vim: ft=python

import sys
import base64
import socket
import xonsh

from prompt_toolkit import print_formatted_text
from prompt_toolkit.formatted_text import PygmentsTokens
from xonsh.pyghooks import partial_color_tokenize

#
# Helpers.
#

iterm2_signal   = lambda contents : f"\033]133;{contents}\a"
iterm2_variable = lambda contents : f"\033]1337;{contents}\a"
iterm2_hostname = socket.gethostname().strip().replace("\n", "")


def printr(raw):
    print(raw, flush=True, end='')

#
def get_iterm2_hostname():
    return iterm2_variable(f"RemoteHost={$USER}@{iterm2_hostname}")

def get_iterm2_cwd():
    return iterm2_variable(f"CurrentDir={$PWD}")

#
# iterm2 integration
#

@events.on_precommand
def iterm2_precmd(cmd=None):
    printr(iterm2_signal('C;'))


@events.on_postcommand
def iterm2_postcmd(cmd, rtn, out, ts):
    printr(iterm2_signal(f"D;{rtn}"))


def iterm2_prompt_formatter(params):

    for index, token in enumerate(params.tokens):
        if token.field == 'iterm2_start':
            prologue = iterm2_signal('A') + get_iterm2_hostname() + get_iterm2_cwd()
            params.update(index, prologue, None, None)

        if token.field == 'iterm2_end':
            epilogue = iterm2_signal('B')
            params.update(index, epilogue, None, None)

    return xonsh.prompt.base.prompt_tokens_formatter_default(params)

#
# Activate shell integration.
#
$SHELL_TYPE = 'readline'
$PROMPT_TOKENS_FORMATTER = iterm2_prompt_formatter

printr(get_iterm2_hostname())
printr(get_iterm2_cwd())
printr(iterm2_variable("ShellIntegrationVersion=16;shell=xonsh"))