#!/usr/bin/env python
#
#    tmux-menu
#    Copyright (C) 2010 Canonical Ltd.
#    Copyright (C) 2012 Dustin Kirkland <kirkland@byobu.co>
#    Copyright (C) 2013 Forkk
#
#    Authors: Dustin Kirkland <kirkland@byobu.co>
#             Ryan C. Thompson <rct@thompsonclan.org>
#             Forkk <forkk@forkk.net>
#
#    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, version 3 of the License.
#
#    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/>.

# This is a heavily modified version of the byobu-select-session script from byobu.

# vim: set expandtab ts=4 sw=4 sts=4:

# Script for listing and choosing tmux sessions.


import os
import re
import sys
import subprocess
import tty
import termios
import select
import string


SHELL = os.getenv("SHELL", "/bin/bash")
HOME = os.getenv("HOME")

UPDATE_ENVVARS = ["DISPLAY", "DBUS_SESSION_BUS_ADDRESS", "SESSION_MANAGER", "GPG_AGENT_INFO", "XDG_SESSION_COOKIE", "XDG_SESSION_PATH", "GNOME_KEYRING_CONTROL", "GNOME_KEYRING_PID", "GPG_AGENT_INFO", "SSH_AUTH_SOCK", "SSH_AGENT_PID"]
# Hacky function to get a char from input without requiring the user to press return.

def getch():
    fd = sys.stdin.fileno()
    old_settings = termios.tcgetattr(fd)
    try:
        mode = termios.tcgetattr(fd)
        mode[tty.LFLAG] = mode[tty.LFLAG] & ~(termios.ECHO | termios.ICANON)
        termios.tcsetattr(fd, termios.TCSAFLUSH, mode)
        char = sys.stdin.read(1)
        return char
    finally:
        termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)

# This causes... Issues...
#try:
#    rw, wl, xl = select.select([fd], [], [], timeout)
#except select.error:
#    return
#if rw:

def get_sessions():
    sessions = []
    text = []
    i = 0
    output = False
    output = subprocess.Popen(["tmux", "list-sessions"], stdout=subprocess.PIPE).communicate()[0]
    output = output.decode(sys.stdout.encoding)
    if output:
        for s in output.splitlines():
            if s:
                text.append("tmux: %s" % s.strip())
                sessions.append("tmux____%s" % s.split(":")[0])
                i += 1
    return (sessions, text)


def update_environment(session):
    _backend, session_name = session.split("____", 2)
    for var in UPDATE_ENVVARS:
        value = os.getenv(var)
        if value:
            cmd = ["tmux", "setenv", "-t", session_name, var, value]
            print("Sending variable: %s" % (cmd, ))
            subprocess.call(cmd, stdout=open(os.devnull, "w"))


def attach_session(session):
    print("Attaching: [%s]\n" % session)
    update_environment(session)
    backend, session_name = session.split("____", 2)
    # must use the binary, not the wrapper!
    os.execvp("tmux", ["", "-2", "attach", "-t", session_name])

def display_menu(options, message="", default=None, max_tries=-1):
    """
    Displays a menu with the given options in the terminal, headed by the given message.
    Waits for the user to select an option, and then returns that option's key.
    If the user selects an invalid option more than ``max_tries`` times, returns None.
    If ``max_tries`` is less than 1, the function will repeat indefinitely until a valid selection is made.
    The ``options`` argument should be a list of two element tuples representing menu entries.
    The first element of each tuple is the key that should be pressed to select that entry, and the second
    element is the text to display on the menu entry.
    """
    # Make sure default is one of the keys in options.
    assert default == None or default in [opt[0] for opt in options]

    # Also make sure all keys are printable non-whitespace characters with a length of one.
    for key, text in options: assert len(key) == 1 and not key.isspace() and key in string.printable

    tries = 0
    # Write the header message if it's valid.
    if message != "":
        sys.stdout.write(message + "\n");

    # Write the list of options.
    for key, text in options:
        sys.stdout.write("\t[%(key)s] - %(text)s\n" % {"key": key, "text": text})

    try:
        while tries < max_tries or max_tries < 0:
            tries += 1

            if default:
                sys.stdout.write("Select an option [%s]: " % default)
            else:
                sys.stdout.write("Select an option: ")
            sys.stdout.flush()

            try:
                # Wait for the user to select an option.
                choice = getch()
            except KeyboardInterrupt:
                return None

            # If the user pressed return and default was set, choose the default.
            if default != None and default != "" and choice in ["\n", "\r"]:
                return default

            if choice in [opt[0] for opt in options]:
                # If the user chose a key that was in the list of options, return the choice.
                return choice
            else:
                # Otherwise, tell the user to try again.
                if choice.isspace() or not choice in string.printable:
                    sys.stdout.write("\rInvalid input. ")
                else:
                    sys.stdout.write("\rInvalid selection: '%s' " % choice);
    finally:
        sys.stdout.write("\n")
        sys.stdout.flush()


# Maps integers to keys. Used as keys for lists of things.
key_num_map = []
for ctr in range(0, 9): key_num_map.append(str(ctr + 1))
key_num_map.append('0')

def main_menu():
    """
    Displays the main menu.
    """
    sessions, sessions_text = get_sessions()

    # Build the main menu.
    menu = [
            ("n", "Create a new tmux session."),
            ("s", "Run a shell without tmux."),
            ("r", "Connect to a remote server via SSH."),
    ]
    for idx, text in enumerate(sessions_text):
        menu.append((key_num_map[idx], text))

    # Show the main menu.
    choice = display_menu(menu, message="Welcome! Please select an option from below.")

    if choice:
        if choice == "n":
            # Create a new session
            os.execvp("tmux", ["", "new-session", SHELL])
        elif choice == "s":
            os.execvp(SHELL, [SHELL])
        elif choice == "r":
            ssh_menu()
        else:
            try:
                # Translate the choice char to an index using the key_num_map list.
                idx = key_num_map.index(choice)
                # Attach to the chosen session; must use the binary, not the wrapper!
                attach_session(sessions[idx])
            except ValueError as e:
                print("Invalid selection! This probably shouldn't happen: %s" % e)
                sys.exit(1)
    else:
        print("Nothing was selected. Exiting.")

def ssh_menu():
    """
    Displays a list of SSH connections available in ~/.ssh_servers
    The file contains a list of connections. One per line in the following format:
    <key> <user>@<address>
    """

    connections = [
        # (<key>, <address>, <user>)
    ]

    try:
        with open(os.path.expanduser("~/.ssh_servers")) as f:
            for line in f.readlines():
                # Parse the line.
                # The key is the first character of the line.
                key = line[0]
                # To get the connection info, we take everything but the first char, strip whitespace, and split it at the @
                cinfo = line[1:].strip().split("@")
                user = cinfo[0]
                addr = cinfo[1]
                # Now add the entry to the connections list.
                connections.append((key, addr, user))
    except FileNotFoundError:
        print("Unable to find the ssh connections list. File ~/.ssh_servers doesn't exist. Going back to the main menu.")
        main_menu()
        return

    # Build the menu
    back_key = "`"
    menu = [
            (back_key, "Go back to the main menu."),
    ]

    for idx, (key, addr, user) in enumerate(connections):
        menu.append((key, "ssh %s@%s" % (user, addr)))

    choice = display_menu(menu, message="Please select an SSH connection.", default=back_key)

    if choice:
        if choice == back_key:
            # Go back to the main menu.
            main_menu()
        else:
            try:
                # Translate the choice char to an index.
                addr, user = [(addr, user) for key, addr, user in connections if key == choice][0]
                print("ssh connecting to %s@%s" % (user, addr))
                os.execvp("ssh", ["", "%s@%s" % (user, addr)])
            except IndexError as e:
                print("Invalid selection! This probably shouldn't happen: %s" % e)
                sys.exit(1)

if __name__ == "__main__":
    main_menu()

