#!/usr/libexec/platform-python
# vim: tabstop=4 shiftwidth=4 softtabstop=4

# Copyright 2015 Lenovo
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import base64
import optparse
import os
import subprocess
import sys
path = os.path.dirname(os.path.realpath(__file__))
path = os.path.realpath(os.path.join(path, '..', 'lib', 'python'))
if path.startswith('/opt'):
    sys.path.append(path)

import confluent.client as client
import confluent.sortutil as sortutil
import confluent.logreader as logreader
import time
import select
import socket
import re
import tty
import termios
import fcntl
import confluent.screensqueeze as sq
try:
    from PIL import Image
except ImportError:
    Image = None

try:
    # sixel is optional, attempt to import but stub out if unavailable
    import io
    import sixel

    class DumbWriter(sixel.SixelWriter):
        def restore_position(self, output):
            return
except ImportError:
    class DumbWriter():
        def draw(self, imgfile):
            sys.stderr.write("PySixel not detected, Sixel format display not supported\n")


confettypath = os.path.join(os.path.dirname(sys.argv[0]), 'confetty')
argparser = optparse.OptionParser(
    usage="Usage: %prog [options] <noderange> [kill][-- [passthroughoptions]]",
    epilog="Command sequences are available while connected to a console, hit "
           "ctrl-'e', then release ctrl, then 'c', then '?' for a full list. "
           "For example, ctrl-'e', then 'c', then '.' will exit the current "
           "console")
argparser.add_option('-t', '--tile', action='store_true', default=False,
                     help='Tile console windows in the terminal')
argparser.add_option('-l', '--log', action='store_true', default=False,
                     help='Enter log replay mode instead of showing a live console')

argparser.add_option('-T', '--Timestamp', action='store_true', default=False,
                     help= 'Dump log in stdout with timestamps')

argparser.add_option('-s', '--screenshot', action='store_true', default=False,
                     help='Attempt to grab screenshot and render using kitty image protocol')
argparser.add_option('-i', '--interval', type='float',
                     help='Interval in seconds to redraw the screenshot. Currently only '
                          'works for one node')
argparser.add_option('-w','--windowed', action='store_true', default=False,
                     help='Open terminal windows for each node.  The '
                          'environment variable NODECONSOLE_WINDOWED_COMMAND '
                          'should be set, which should be a text string corresponding '
                          'to a command that can be used to open a windowed console,'
                          ' omitting the "nodeconsole <noderange>" part of the '
                          'command, for example, to open a set of consoles for a '
                          'range of nodes in separate xterm windows, set '
                          'NODECONSOLE_WINDOWED_COMMAND to "xterm -e".  To open a '
                          'set of consoles for a range of nodes in separate '
                          'GNOME Terminal windows with a size of 100 columns and '
                          '31 rows, set NODECONSOLE_WINDOWED_COMMAND '
                          'to "gnome-terminal --geometry 100x31 --" or in a WSL '
                          'environment, to open a set of consoles for a range of '
                          'nodes in separate Windows Terminal windows, with the '
                          'title set for each node, set NODECONSOLE_WINDOWED_COMMAND'
                          ' to "wt.exe wsl.exe -d AlmaLinux-8 '
                          '--shell-type login".  If the NODECONSOLE_WINDOWED_COMMAND '
                          'environment variable isn\'t set, xterm will be used by'
                          'default.')

(options, args) = argparser.parse_args()

oldtcattr = None
oldfl = None

def get_coords():
    sys.stdout.write('\x1b[6n') #
    sys.stdout.flush()
    gotreply = select.select([sys.stdin,], [], [], 0.250)[0]
    if gotreply:
        response = ''
        while select.select([sys.stdin,], [], [], 0.1)[0] and 'R' not in response:
            response += sys.stdin.read()
        coords = response.replace('R', '').split('[')[1].split(';')
    #sys.stdout.write('\x1b[{}:{}H'.format(*coords))

def direct_console():
    global oldtcattr
    global oldfl
    oldtcattr = termios.tcgetattr(sys.stdin.fileno())
    oldfl = fcntl.fcntl(sys.stdin.fileno(), fcntl.F_GETFL)
    tty.setraw(sys.stdin.fileno())
    fcntl.fcntl(sys.stdin.fileno(), fcntl.F_SETFL, oldfl | os.O_NONBLOCK)

def indirect_console():
    fcntl.fcntl(sys.stdin.fileno(), fcntl.F_SETFL, oldfl & ~os.O_NONBLOCK)
    termios.tcsetattr(sys.stdin.fileno(), termios.TCSANOW, oldtcattr)

def determine_tile_size(numnodes):
    # for now, smash everything to a common aspect ratio.  16:11
    # is pretty much wrong for everything, making 4:3 a bit too wide
    # and 16:9 significantly too narrow, but it is serviceable
    # An improvement could come with us owning the scaling
    # instead of delegating to Kitty, which says if we specify both,
    # we get stretching.  In theory we should be able to get aspect correct
    # from kitty by omitting, but:
    # then we don't know how much to move the cursor left after draw_image
    # Konsole won't scale at all with only partial scaling specified
    cheight, cwidth, pixwidth, pixheight = sq.get_screengeom()
    # 16:12 is to roughly account for the 'titles' of the tiles
    ratio = (pixwidth / 16) / (pixheight / 12)
    bestdeviation = None
    bestdims = []
    for i in range(1, numnodes + 1):
        number = numnodes
        while number % i != 0:
            number += 1
        columns = i
        rows = number // i
        deviation = abs(ratio -  (columns / rows))
        if bestdeviation is None:
            bestdeviation = deviation
            bestdims = [columns, rows]
        elif deviation < bestdeviation:
            bestdeviation = deviation
            bestdims = [columns, rows]
    # ok, the above algorithm can still pick things like
    # 1 2 3
    # 4
    # So we will let it pick the number of rows, and
    # then see if we can chop columns and still fit
    while (bestdims[0] - 1) * bestdims[1] >= numnodes:
        bestdims[0] = bestdims[0] - 1
    cellswide = cwidth // bestdims[0]
    cellshigh = cheight // bestdims[1]
    tilewidth = cellswide * pixwidth / cwidth
    tileheight = cellshigh * pixheight / cheight
    if tilewidth > (tileheight * 16 / 11):
        tilewidth = tileheight * 16 / 11
        cellswide = int(tilewidth // (pixwidth / cwidth))
    if tileheight > (tilewidth * 11 /16):
        tileheight = tilewidth * 11 / 16
        cellshigh = int(tileheight // (pixheight / cheight))
    bestdims = bestdims + [cellswide, cellshigh, cellshigh * bestdims[1]]
    # incur any scrolling we might get. This allows us to accurately
    # save/restore cursor or even get coordinates without scrolling fouling
    # the desired target
    sys.stdout.write('\n' * bestdims[4])
    sys.stdout.flush()
    cursor_up(bestdims[4])
    return bestdims

cursor_saved = False
def sticky_cursor():
    global cursor_saved
# get cursor restore_position
    if sys.stdin.isatty() and not cursor_saved:
        try:
            direct_console()
            sys.stdout.write('\x1b7')
            cursor_saved = True
        finally:
            indirect_console()
    elif cursor_saved:
        try:
            direct_console()
            sys.stdout.write('\x1b8')
        finally:
            indirect_console()

def cursor_up(count=1):
    sys.stdout.write(f'\x1b[{count}A')
def cursor_down(count=1):
    sys.stdout.write(f'\x1b[{count}B')
def cursor_right(count=1):
    sys.stdout.write(f'\x1b[{count}C')
def cursor_left(count=1):
    sys.stdout.write(f'\x1b[{count}D')
def cursor_save():
    sys.stdout.write('\x1b7')
def cursor_restore():
    sys.stdout.write('\x1b8')
def cursor_hide():
    sys.stdout.write('\x1b[?25l')
def cursor_show():
    sys.stdout.write('\x1b[?25h')

def draw_image(data, width, height):
    imageformat = os.environ.get('CONFLUENT_IMAGE_PROTOCOL', 'kitty')
    if imageformat == 'sixel':
        sixel_draw(data)
    elif imageformat == 'iterm':
        iterm_draw(data, width, height)
    else:
        kitty_draw(data, width, height)


def sixel_draw(data):
    bindata = base64.b64decode(data)
    binfile = io.BytesIO()
    binfile.write(bindata)
    binfile.seek(0)
    DumbWriter().draw(binfile)

def iterm_draw(data, width, height):
    if not height:
        height = 'auto'
    if not width:
        width = 'auto'
    bindata = base64.b64decode(data)
    datalen = len(bindata)
    sys.stdout.write(
        '\x1b]1337;File=inline=1;width={};height={};size={}:'.format(width,height,datalen))
    sys.stdout.write(data.decode('utf8'))
    sys.stdout.write('\a')
    sys.stdout.flush()

def kitty_draw(data, width, height):
    if Image:
        bindata = base64.b64decode(data)
        binfile = io.BytesIO()
        binfile.write(bindata)
        binfile.seek(0)
        img = Image.open(binfile)
        outfile = io.BytesIO()
        img.save(outfile, format='PNG')
        data = base64.b64encode(outfile.getbuffer())
    preamble = '\x1b_Ga=T,f=100'
    if height:
        preamble += f',r={height},c={width}'
    #sys.stdout.write(repr(preamble))
    #sys.stdout.write('\xb[{}D'.format(len(repr(preamble))))
    #return
    first = True
    while data:
        chunk, data = data[:4096], data[4096:]
        m = 1 if data else 0
        if first:
            sys.stdout.write('{},m={};'.format(preamble, m))
        else:
            sys.stdout.write('\x1b_Gm={};'.format(m))
        sys.stdout.write(chunk.decode('utf8'))
        sys.stdout.write('\x1b\\')
        sys.stdout.flush()

pass_through_args = []
killcon = False
try:
    noderange = args[0]
    if len(args) > 1:
        if args[1] == 'kill':
            killcon = True
        pass_through_args = args[1:]
        args = args[:1]
except IndexError:
    argparser.print_help()
    sys.exit(1)

if len(args) != 1:
    argparser.print_help()
    sys.exit(1)

if options.log:
    logname = args[0]
    if not os.path.exists(logname) and logname[0] != '/':
        logname = os.path.join('/var/log/confluent/consoles', logname)
    if not os.path.exists(logname):
        sys.stderr.write('Unable to locate {0} on local system\n'.format(logname))
        sys.exit(1)
    logreader.replay_to_console(logname)
    sys.exit(0)

if options.Timestamp:
    logname = args[0]
    if not os.path.exists(logname) and logname[0] != '/':
        logname = os.path.join('/var/log/confluent/consoles', logname)
    if not os.path.exists(logname):
        sys.stderr.write('Unable to locate {0} on local system\n'.format(logname))
        sys.exit(1)
    logreader.dump_to_console(logname)
    sys.exit(0)

def prep_node_tile(node):
    currcolcell, currrowcell = nodepositions[node]
    if currcolcell:
        cursor_right(currcolcell)
    if currrowcell:
        cursor_down(currrowcell)
    sys.stdout.write(node)
    cursor_left(len(node))
    cursor_down()

def reset_cursor(node):
    currcolcell, currrowcell = nodepositions[node]
    if currcolcell:
        cursor_left(currcolcell)
    cursor_up(currrowcell + 1)

nodepositions = {}
numrows = 0
def do_screenshot():
    global numrows
    cwidth = None
    cheight = None
    sess = client.Command()
    if options.tile:
        imageformat = os.environ.get('CONFLUENT_IMAGE_PROTOCOL', 'kitty')
        if imageformat not in ('kitty', 'iterm'):
            sys.stderr.write('Tiled screenshots only supported with kitty or iterm protocol')
            sys.exit(1)
        allnodes = []
        numnodes = 0
        for res in sess.read('/noderange/{}/nodes/'.format(args[0])):
            allnodes.append(res['item']['href'].replace('/', ''))
            numnodes += 1
        cols, rows, cwidth, cheight, numrows = determine_tile_size(numnodes)
        currcol = 1
        currcolcell = 0
        currrowcell = 0
        for node in allnodes:
            nodepositions[node] = currcolcell, currrowcell
            if currcol < cols:
                currcol += 1
                currcolcell += cwidth
            else:
                currcol = 1
                currcolcell = 0
                currrowcell += cheight
    elif options.interval is not None:
        sys.stdout.write('\x1bc')
    firstnodename = None
    dorefresh = True
    while dorefresh:
        for res in sess.read('/noderange/{}/console/ikvm_screenshot'.format(args[0])):
            for node in res.get('databynode', {}):
                if not firstnodename:
                    firstnodename = node
                imgdata = res['databynode'][node].get('image', {}).get('imgdata', None)
                if imgdata:
                    if len(imgdata) < 32: # We were subjected to error
                        sys.stderr.write(f'{node}: Unable to get screenshot\n')
                        continue
                    if node in nodepositions:
                        prep_node_tile(node)
                        cursor_save()
                    else:
                        if options.interval is not None:
                            if node != firstnodename:
                                sys.stderr.write('Multiple nodes not supported for interval')
                                sys.exit(1)
                            sticky_cursor()
                        sys.stdout.write('{}: '.format(node))
                    # one row is used by our own name, so cheight - 1 for that allowance
                    draw_image(imgdata.encode(), cwidth, cheight - 1 if cheight else cheight)
                    if node in nodepositions:
                        cursor_restore()
                        reset_cursor(node)
                    else:
                        sys.stdout.write('\n')
                    sys.stdout.flush()
        if options.interval is None:
            dorefresh = False
        else:
            dorefresh = True
            time.sleep(options.interval)
    sys.exit(0)

if options.screenshot:
    try:
        cursor_hide()
        do_screenshot()
    except KeyboardInterrupt:
        pass
    finally:
        cursor_show()
        cursor_down(numrows)
        sys.stdout.write('\n')
    sys.exit(0)
def kill(noderange):
    sess = client.Command()
    envstring=os.environ.get('NODECONSOLE_WINDOWED_COMMAND')
    if not envstring:
        envstring = 'xterm'

    nodes = []
    for res in sess.read('/noderange/{0}/nodes/'.format(args[0])):
        node = res.get('item', {}).get('href', '/').replace('/', '')
        if not node:
            sys.stderr.write(res.get('error', repr(res)) + '\n')
            sys.exit(1)
        nodes.append(node) 

    for node in nodes:
        s=socket.socket(socket.AF_UNIX)
        winid=None
        try:
            win=subprocess.Popen(['xwininfo', '-tree', '-root'], stdout=subprocess.PIPE)
            wintr=win.communicate()[0]
            for line in wintr.decode('utf-8').split('\n'):
                if 'console: {0}'.format(node) in line or 'confetty' in line:
                    win_obj = [ele for ele in line.split(' ') if ele.strip()]
                    winid = win_obj[0]
        except:
            print("Error: cannot retrieve window id of node {}".format(node))

        if winid:
            ps_data=subprocess.Popen(['xkill', '-id', winid ], stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    sys.exit(0)

def handle_geometry(envlist, sizegeometry, side_pad=0, top_pad=0, first=False):    
    if '-geometry' in envlist:
        g_index = envlist.index('-geometry')
    elif '-g' in envlist:
        g_index = envlist.index('-g')
    else:
        g_index = 0
    if g_index:   
        if first:
             envlist[g_index+1] = '{0}+{1}+{2}'.format(envlist[g_index+1],side_pad, top_pad)
        else:
            envlist[g_index+1] = '{0}+{1}+{2}'.format(sizegeometry,side_pad, top_pad)
    else:
        envlist.insert(1, '-geometry')
        envlist.insert(2, '{0}+{1}+{2}'.format(sizegeometry,side_pad, top_pad))
        g_index = 1
    return envlist     

# add funcltionality to close/kill all open consoles
if killcon:
    kill(noderange)
    
#added functionality for wcons
if options.windowed:
    result=subprocess.Popen(['xwininfo', '-root'], stdout=subprocess.PIPE)
    rootinfo=result.communicate()[0]
    result.wait()
    for line in rootinfo.decode('utf-8').split('\n'):
        if 'Width' in line:
            screenwidth = int(line.split(':')[1])
        if 'Height' in line:
            screenheight = int(line.split(':')[1])

    envstring=os.environ.get('NODECONSOLE_WINDOWED_COMMAND')
    if not envstring:
        sizegeometry='100x31'
        corrected_x, corrected_y = (13,84)
        envlist = handle_geometry(['xterm'] + pass_through_args + ['-e'],sizegeometry, first=True)
        #envlist=['xterm', '-bg', 'black', '-fg', 'white', '-geometry', '{sizegeometry}+0+0'.format(sizegeometry=sizegeometry), '-e']
    else:
        envlist=os.environ.get('NODECONSOLE_WINDOWED_COMMAND').split(' ')
        if envlist[0] == 'xterm':
            if '-geometry' in envlist:
                g_index = envlist.index('-geometry')
            elif '-g' in envlist:
                g_index = envlist.index('-g')
            else:
                g_index = 0
            if g_index:    
                envlist[g_index+1] = envlist[g_index+1] + '+0+0'

            else:
                envlist.insert(1, '-geometry')
                envlist.insert(2, '100x31+0+0')
                g_index = 1
            
    nodes = []
    sess = client.Command()
    for res in sess.read('/noderange/{0}/nodes/'.format(args[0])):
        node = res.get('item', {}).get('href', '/').replace('/', '')
        if not node:
            sys.stderr.write(res.get('error', repr(res)) + '\n')
            sys.exit(1)
        nodes.append(node)

    if options.tile and not envlist[0] == 'xterm':
        sys.stderr.write('[ERROR] UNSUPPORTED OPTIONS. \nWindowed and tiled consoles are only supported when using xterm \n')
        sys.exit(1)
    firstnode=nodes[0]
    nodes.pop(0)
    with open(os.devnull, 'wb') as devnull:
        xopen=subprocess.Popen(envlist + [confettypath, '-c', '/tmp/controlpath-{0}'.format(firstnode), '-m', '5', 'start', '/nodes/{0}/console/session'.format(firstnode) ] , stdin=devnull)
        time.sleep(2)
        s=socket.socket(socket.AF_UNIX)
        winid=''
        try:
            s.connect('/tmp/controlpath-{firstnode}'.format(firstnode=firstnode))
            s.recv(64)
            s.send(b'GETWINID')
            winid=s.recv(64).decode('utf-8')

        except:
            time.sleep(2)
            # try to get id of first panel/xterm window using name
            win=subprocess.Popen(['xwininfo', '-tree', '-root'], stdout=subprocess.PIPE)
            wintr=win.communicate()[0]
            for line in wintr.decode('utf-8').split('\n'):
                if 'console: {firstnode}'.format(firstnode=firstnode) in line or 'confetty' in line:
                    win_obj = [ele for ele in line.split(' ') if ele.strip()]
                    winid = win_obj[0]
    if winid:
        firstnode_window=subprocess.Popen(['xwininfo', '-id', '{winid}'.format(winid=winid)], stdout=subprocess.PIPE)
        xinfo=firstnode_window.communicate()[0]
        xinfl = xinfo.decode('utf-8').split('\n')
        for line in xinfl:
            if 'Absolute upper-left X:' in line:
                side_pad = int(line.split(':')[1])
            elif 'Absolute upper-left Y:' in line:
                top_pad = int(line.split(':')[1])

            elif 'Width:' in line:
                window_width = int(line.split(':')[1])
            elif 'Height' in line:
                window_height = int(line.split(':')[1])
            elif '-geometry' in line:
                l = re.split(' |x|-|\\+', line)
                l_nosp = [ele for ele in l if ele.strip()]
                wmxo = int(l_nosp[1])
                wmyo = int(l_nosp[2])
                sizegeometry = str(wmxo) + 'x' + str(wmyo)
            else:
                pass

        window_width += side_pad*2
        window_height += side_pad+top_pad
        screenwidth -= wmxo
        screenheight -= wmyo
        currx = window_width
        curry = 0
        maxcol = int(screenwidth/window_width)

    for node in sortutil.natural_sort(nodes):
        if options.tile and envlist[0] == 'xterm':
            corrected_x = currx
            corrected_y = curry
            xgeometry = '{0}+{1}+{2}'.format(sizegeometry, corrected_x, corrected_y)
            currx += window_width
            if currx >= screenwidth:
                currx=0
                curry += window_height
                if curry > screenheight:
                    curry =top_pad
            if not envstring:        
                envlist= handle_geometry(envlist, sizegeometry, corrected_x, corrected_y)
            else:
                if g_index:    
                    envlist[g_index+1] = xgeometry
        elif envlist[0] == 'xterm':
            envlist=handle_geometry(envlist, sizegeometry, side_pad, top_pad)
            side_pad+=(side_pad+1)
            top_pad+=(top_pad+30)
        else:
            pass 
        with open(os.devnull, 'wb') as devnull:
            xopen=subprocess.Popen(envlist + [confettypath, '-c', '/tmp/controlpath-{0}'.format(node), '-m', '5', 'start', '/nodes/{0}/console/session'.format(node)] , stdin=devnull)
    sys.exit(0)
#end of wcons
if options.tile:
    null = open('/dev/null', 'w')
    nodes = []
    sess = client.Command()
    for res in sess.read('/noderange/{0}/nodes/'.format(args[0])):
        node = res.get('item', {}).get('href', '/').replace('/', '')
        if not node:
            sys.stderr.write(res.get('error', repr(res)) + '\n')
            sys.exit(1)
        nodes.append(node)
    initial = True
    in_tmux = False
    pane = 0
    sessname = 'nodeconsole_{0}'.format(os.getpid())
    if os.environ.get("TMUX"):
        initial = False
        in_tmux = True
        subprocess.call(['tmux', 'rename-session', sessname])
    for node in sortutil.natural_sort(nodes):
        panename = '{0}:{1}'.format(sessname, pane)
        if initial:
            initial = False
            subprocess.call(
                ['tmux', 'new-session', '-d', '-s',
                 sessname, '-x', '800', '-y',
                 '800', '{0} -m 5 start /nodes/{1}/console/session'.format(
                    confettypath, node)])
        else:
            subprocess.call(['tmux', 'select-pane', '-t', sessname])
            subprocess.call(['tmux', 'set-option', '-t', panename, 'pane-border-status', 'top'], stderr=null)
            subprocess.call(
                    ['tmux', 'split', '-h', '-t', sessname,
                 '{0} -m 5 start /nodes/{1}/console/session'.format(
                     confettypath, node)])
            subprocess.call(['tmux', 'select-layout', '-t', sessname, 'tiled'], stdout=null)
            pane += 1
    subprocess.call(['tmux', 'select-pane', '-t', sessname])
    subprocess.call(['tmux', 'set-option', '-t', panename, 'pane-border-status', 'top'], stderr=null)
    if not in_tmux:
        os.execlp('tmux', 'tmux', 'attach', '-t', sessname)
else:
    os.execl(confettypath, confettypath, 'start',
             '/nodes/{0}/console/session'.format(args[0]))