#!/usr/bin/env python

import argparse
import base64
import json
import os
import sys
import uuid
from getpass import getpass
from subprocess import check_output
from enum import Enum

import databrickslabs_jupyterlab
from databrickslabs_jupyterlab.remote import (bye, check_installed, configure_ssh, connect, get_cluster,
                                              get_remote_packages, install_libs, is_reachable, mount_sshfs,
                                              version_check)
from databrickslabs_jupyterlab.local import (bye, conda_version, create_kernelspec, get_db_config, prepare_ssh_config,
                                             print_error, print_ok, show_profiles, write_config)
from databrickslabs_jupyterlab.install import install, update_local


class VCheck(Enum):
    a = 'all'
    d = 'diff'
    s = 'same'

    def __str__(self):
        return self.value

parser = argparse.ArgumentParser(description='Configure remote Databricks access with Jupyter Lab')

parser.add_argument('profile', nargs='?', type=str, default="", help='A databricks-cli profile')
parser.add_argument('-b',
                    '--bootstrap',
                    dest='bootstrap',
                    action='store_true',
                    help='Bootstrap the local databrickslabs-jupyterlab environment')
parser.add_argument('-m',
                    '--mirror',
                    dest='mirror',
                    action='store_true',
                    help='Mirror a a remote Databricks environment')
parser.add_argument('-c',
                    '--clipboard',
                    dest='clipboard',
                    action='store_true',
                    help='Copy the personal access token to the clipboard')
parser.add_argument('-f',
                    '--force',
                    dest='force',
                    action='store_true',
                    help='Force remote installation of databrickslabs_jupyterlab package')
parser.add_argument('-i', '--id', dest='cluster_id', help='The cluster_id to avoid manual selection')
parser.add_argument('-k', '--kernelspec', dest='kernelspec', action='store_true', help='Create a kernel specification')
parser.add_argument('-l', '--lab', dest='lab', action='store_true', help='Safely start Jupyter Lab')
parser.add_argument('-o', '--organisation', dest='organisation', help='The organisation for Azure Databricks')
parser.add_argument('-p',
                    '--profiles',
                    dest='profiles',
                    action='store_true',
                    help='Show all databricks cli profiles and check SSH key')
parser.add_argument('-r',
                    '--reconfigure',
                    dest='reconfigure',
                    action='store_true',
                    help='Reconfigure cluster with id cluster_id')
# parser.add_argument('-s', '--sshfs', dest='sshfs', action='store_true', help='Mount remote filesystem via sshfs')
parser.add_argument('-s',
                    '--ssh-config',
                    dest='sshconfig',
                    action='store_true',
                    help='Configure SSH acces for a cluster')
parser.add_argument('-v',
                    '--version',
                    dest='version',
                    action='store_true',
                    help='Check version of databrickslabs-jupyterlab')
parser.add_argument('-V',
                    '--versioncheck',
                    dest='versioncheck',
                    type=VCheck,
                    choices=list(VCheck),
                    help='Check version of local env with remote env')

args = parser.parse_args()

#
# Print Version and exit
#
if args.version:
    import databrickslabs_jupyterlab._version
    print(databrickslabs_jupyterlab._version.__version__)
    sys.exit(0)

#
# Bootstrap the initial environment and exit
#
if args.bootstrap:
    update_local()
    sys.exit(0)

#
# Show all configure profiles
#
if args.profiles:
    show_profiles()
    exit(0)


profile = args.profile

#
# Check that at least conda 4.7.5 is installed. If not, exit
#
version = conda_version()
if version[1] >= 4 and version[2] >= 5:
    print("Valid version of conda detected: %s.%s.%s" % version)
else:
    print("Too old conda version:")
    print("Please update conda to at least 4.7.5")
    sys.exit(1)

#
# Ensure profile argument is given
#
if args.profile == "":
    parser.print_help()
    print_error("\ndatabrickslabs-jupyterlab: error: the following arguments are required: profile\n")
    bye(1)

print("\n* Getting host and token from .databrickscfg")
host, token = get_db_config(profile)
if ("azuredatabricks.net" in host) and args.kernelspec and not args.organisation:
    print_error("\n    Error: To configure an Azure Databricks cluster, " +
                "the organization id (-o id) is required")
    bye(1)

#
# Configure SSH for cluster
#
if args.sshconfig:
    if not args.cluster_id:
        print_error("   => Error: Provide '-i cluster_id' with '-s' ")
        bye(1)
    configure_ssh(profile, host, token, args.cluster_id)
    bye()

#
# Check the ssh key for the given profile and retrieve host and token from ~/.databrickscfg
#
if not os.path.exists(os.path.expanduser("~/.ssh/id_%s.pub" % profile)):
    print_error("\n    Error: ssh key for profile '%s' is missing:" %profile)
    print("    Use 'databrickslabs-jupyterlab %s -s' to configure ssh access" % profile)
    bye(1)


#
# Copy PAT to clipboard. Can be combined with every other flag
#
if args.clipboard:
    import pyperclip
    pyperclip.copy(token)
    print_ok("   => Personal access token copied to clipboard")

#
# Select remote cluster if not given by either -i or being in the respective mirror environment
#
if args.kernelspec or args.reconfigure or args.lab or args.versioncheck or args.mirror:
    print("\n* Select remote cluster")
    cluster_id, public_ip, cluster_name, started = get_cluster(profile, host, token, args.cluster_id)
    if cluster_name is None:
        bye(1)

    print("\n* Configuring ssh config for remote cluster")
    prepare_ssh_config(cluster_id, profile, public_ip)
    print("   => Testing whether cluster can be reached")
    if not is_reachable(public_dns=public_ip):
        print_error("Cannot connect to remote cluster. Please check:")
        print_error("- whether port 2200 is open in cloud security group")
        print_error("- whether VPN is enabled if you need one to connect to the cluster")
        bye(1)

#
# Compare local and remote libray versions
#
if args.versioncheck:
    version_check(cluster_id, args.versioncheck)
    bye()


#
# Create a conda environment that mirrors the selected remote cluster and exit
#
if args.mirror:
    install(profile, cluster_id, cluster_name)
    bye()

#
# Install databrickslabs_jupyterlab libraries on the driver
#
if args.kernelspec or args.reconfigure or args.lab:
    print("   => Check installed libraries")
    if args.force or not check_installed(cluster_id):
        print("\n* Installing databrickslabs_jupyterlab, ipywidgets on remote driver")
        install_libs(cluster_id)

    # Recheck in case something went wrong
    if check_installed(cluster_id):
        print_ok("   => OK")
    else:
        print_error("databrickslabs_jupyterlabs is not installed on the remote cluster")
        bye(1)

#
# Start Jupyter lab in a safely way with all prechecks from above
#
if args.lab:
    import jupyterlab.labapp
    sys.argv = sys.argv[:1]
    print(sys.argv)
    jupyterlab.labapp.main()

#
# Create a jupyter kernelspecification
#
if args.kernelspec:
    print("\n* Create remote kernel spec")
    # Using local conda environment name avoids overwriting of kernelspecs from different environments
    create_kernelspec(profile, args.organisation, host, cluster_id, os.environ['CONDA_DEFAULT_ENV'])

    print("\n* Setting global config of jupyter lab (autorestart, timeout)")
    write_config()

print("")
