#!python

import argparse
import os
import platform
import random
import shutil
import string
import subprocess
import sys
import yaml

import distro

base_directory = os.path.join(os.environ["HOME"], ".aqueduct")
server_directory = os.path.join(os.environ["HOME"], ".aqueduct", "server")
ui_directory = os.path.join(os.environ["HOME"], ".aqueduct", "ui")


def update_config_yaml(file):
    s=string.ascii_uppercase+string.digits
    encryption_key = ''.join(random.sample(s,32))
    api_key = ''.join(random.sample(s,32))

    with open(file, "r") as sources:
        lines = sources.readlines()
    with open(file, "w") as sources:
        for line in lines:
            if "<BASE_PATH>" in line:
                sources.write(line.replace("<BASE_PATH>", server_directory))
            elif "<ENCRYPTION_KEY>" in line:
                sources.write(line.replace("<ENCRYPTION_KEY>", encryption_key))
            elif "<API_KEY>" in line:
                sources.write(line.replace("<API_KEY>", api_key))
            else:
                sources.write(line)

def generate_env_local(ip):
    if os.path.isfile(os.path.join(ui_directory, "app", ".env.local")):
        execute_command(["rm", os.path.join(ui_directory, "app", ".env.local")])

    execute_command(["touch", os.path.join(ui_directory, "app", ".env.local")])

    with open(os.path.join(ui_directory, "app", ".env.local"), "w") as env_file:
        env_file.write("SERVER_ADDRESS=%s:8080\n" % ip)
        env_file.write("NEXT_PUBLIC_PROTOCOL=http\n")

def execute_command(args, cwd=None):
    with subprocess.Popen(args, stdout=sys.stdout, stderr=sys.stderr, cwd=cwd) as proc:
        proc.communicate()
        if proc.returncode != 0:
            raise Exception("Error executing command: %s" % args)

def update_executable_permissions():
    execute_command(["chmod", "755", os.path.join(server_directory, "bin", "server")])
    execute_command(["chmod", "755", os.path.join(server_directory, "bin", "executor")])
    execute_command(["chmod", "755", os.path.join(server_directory, "bin", "migrator")])

def server():
    if not os.path.isdir(server_directory):
        try:
            directories = [
                server_directory,
                os.path.join(server_directory, "db"),
                os.path.join(server_directory, "storage"),
                os.path.join(server_directory, "storage", "operators"),
                os.path.join(server_directory, "vault"),
                os.path.join(server_directory, "bin"),
                os.path.join(server_directory, "config"),
            ]

            for directory in directories:
                os.mkdir(directory)

            system = platform.system()
            arch = platform.machine()
            if system == "Linux" and arch == "x86_64":
                print("Operating system is Linux with architecture amd64.")
                execute_command(["curl", "https://aqueduct-ai.s3.us-east-2.amazonaws.com/assets/bin/linux_amd64/server", "--output", os.path.join(server_directory, "bin", "server")])
                execute_command(["curl", "https://aqueduct-ai.s3.us-east-2.amazonaws.com/assets/bin/linux_amd64/executor", "--output", os.path.join(server_directory, "bin", "executor")])
                execute_command(["curl", "https://aqueduct-ai.s3.us-east-2.amazonaws.com/assets/bin/linux_amd64/migrator", "--output", os.path.join(server_directory, "bin", "migrator")])
            elif system == "Darwin" and arch == "x86_64":
                print("Operating system is Mac with architecture amd64.")
                execute_command(["curl", "https://aqueduct-ai.s3.us-east-2.amazonaws.com/assets/bin/darwin_amd64/server", "--output", os.path.join(server_directory, "bin", "server")])
                execute_command(["curl", "https://aqueduct-ai.s3.us-east-2.amazonaws.com/assets/bin/darwin_amd64/executor", "--output", os.path.join(server_directory, "bin", "executor")])
                execute_command(["curl", "https://aqueduct-ai.s3.us-east-2.amazonaws.com/assets/bin/darwin_amd64/migrator", "--output", os.path.join(server_directory, "bin", "migrator")])
            elif system == "Darwin" and arch == "arm64":
                print("Operating system is Mac with architecture arm64.")
                execute_command(["curl", "https://aqueduct-ai.s3.us-east-2.amazonaws.com/assets/bin/darwin_arm64/server", "--output", os.path.join(server_directory, "bin", "server")])
                execute_command(["curl", "https://aqueduct-ai.s3.us-east-2.amazonaws.com/assets/bin/darwin_arm64/executor", "--output", os.path.join(server_directory, "bin", "executor")])
                execute_command(["curl", "https://aqueduct-ai.s3.us-east-2.amazonaws.com/assets/bin/darwin_arm64/migrator", "--output", os.path.join(server_directory, "bin", "migrator")])
            else:
                raise Exception("Unsupported operating system and architecture combination: %s, %s" % (system, arch))
            
            update_executable_permissions()

            execute_command(["curl", "https://aqueduct-ai.s3.us-east-2.amazonaws.com/assets/bin/start-function-executor.sh", "--output", os.path.join(server_directory, "bin", "start-function-executor.sh")])
            execute_command(["curl", "https://aqueduct-ai.s3.us-east-2.amazonaws.com/assets/bin/install_sqlserver_ubuntu.sh", "--output", os.path.join(server_directory, "bin", "install_sqlserver_ubuntu.sh")])
            execute_command(["curl", "https://aqueduct-ai.s3.us-east-2.amazonaws.com/assets/db/demo.db", "--output", os.path.join(server_directory, "db", "demo.db")])
            execute_command(["curl", "https://aqueduct-ai.s3.us-east-2.amazonaws.com/assets/config/config.yml", "--output", os.path.join(server_directory, "config", "config.yml")])

            update_config_yaml(os.path.join(server_directory, "config", "config.yml"))

            execute_command([os.path.join(server_directory, "bin", "migrator"), "--type", "sqlite", "goto", "8"])

            print("Finished initializing Aqueduct base directory.")
        except Exception as e:
            print(e)
            execute_command(["rm", "-rf", server_directory])
            exit(1)

    execute_command([os.path.join(server_directory, "bin", "server"), "--config", os.path.join(server_directory, "config", "config.yml")])

def install_postgres():
    execute_command([sys.executable, "-m", "pip", "install", "psycopg2-binary"])

def install_bigquery():
    execute_command([sys.executable, "-m", "pip", "install", "pandas-gbq"])

def install_snowflake():
    execute_command([sys.executable, "-m", "pip", "install", "snowflake-sqlalchemy"])

def install_s3():
    execute_command([sys.executable, "-m", "pip", "install", "pyarrow"])

def install_mysql():
    system = platform.system()
    if system == "Linux":
        if distro.id() == "ubuntu" or distro.id() == "debian":
            execute_command(["sudo", "apt-get", "install", "-y", "python3-dev", "default-libmysqlclient-dev", "build-essential"])
        elif distro.id() == "centos" or distro.id() == "rhel":
            execute_command(["sudo", "yum", "install", "-y", "python3-devel", "mysql-devel"])
        else:
            print("Unsupported distribution:", distro.id())
    elif system == "Darwin":
        execute_command(["brew", "install", "mysql"])
    else:
        print("Unsupported operating system:", system)
    
    execute_command([sys.executable, "-m", "pip", "install", "mysqlclient"])

def install_sqlserver():
    system = platform.system()
    if system == "Linux":
        if distro.id() == "ubuntu":
            execute_command(["bash", os.path.join(server_directory, "bin", "install_sqlserver_ubuntu.sh")])
        else:
            print("Unsupported distribution:", distro.id())
    elif system == "Darwin":
        execute_command(["brew", "tap", "microsoft/mssql-release", "https://github.com/Microsoft/homebrew-mssql-release"])
        execute_command(["brew", "update"])
        execute_command(["HOMEBREW_NO_ENV_FILTERING=1", "ACCEPT_EULA=Y", "brew", "install", "msodbcsql17", "mssql-tools"])
    else:
        print("Unsupported operating system:", system)
    
    execute_command([sys.executable, "-m", "pip", "install", "pyodbc"])

def install(system):
    if system == "postgres":
        install_postgres()
    elif system == "bigquery":
        install_bigquery()
    elif system == "snowflake":
        install_snowflake()
    elif system == "s3":
        install_s3()
    elif system == "mysql":
        install_mysql()
    elif system == "sqlserver":
        install_sqlserver()
    else:
        raise Exception("Unsupported system: %s" % system)

def ui(exposeIp):
    if not os.path.isdir(ui_directory):
        try:
            os.mkdir(ui_directory)
            execute_command(["curl", "https://aqueduct-ai.s3.us-east-2.amazonaws.com/assets/ui/ui.zip", "--output", os.path.join(ui_directory, "ui.zip")])
            execute_command(["unzip", os.path.join(ui_directory, "ui.zip"), "-d", ui_directory])
            execute_command(["rm", os.path.join(ui_directory, "ui.zip")])
            execute_command(["npm", "install"], os.path.join(ui_directory, "app"))
            execute_command(["npm", "run", "build"], os.path.join(ui_directory, "app"))
        except Exception as e:
            print(e)
            execute_command(["rm", "-rf", ui_directory])
            exit(1)

    if exposeIp:
        server_ip = exposeIp
        print("The Aqueduct server at " + server_ip + " will be externally accessible.")
    else:
        server_ip = "localhost"

    generate_env_local(server_ip)
    execute_command(["npm", "run", "start"], os.path.join(ui_directory, "app"))

def apiKey():
    config_file = os.path.join(server_directory, "config", "config.yml")
    with open(config_file, "r") as f:
        try:
            print(yaml.safe_load(f)['apiKey'])
        except yaml.YAMLError as exc:
            print(exc)
            exit(1)

def clear():
    execute_command(["rm", "-rf", base_directory])

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='The Aqueduct CLI')
    subparsers = parser.add_subparsers(dest="command")
    
    server_args = subparsers.add_parser('server', help=
                                   '''This starts the Aqueduct server in a
                                   blocking fashion. To background the process,
                                   run aqueduct server &.''')
    ui_args = subparsers.add_parser('ui', help=
                               '''This starts the Aqueduct UI in a blocking
                               fashion. To background the process run aqueduct
                               ui &.
                               Add --expose <IP_ADDRESS> to access the UI from
                               an external server, where <IP_ADDRESS> is the
                               public IP of the server you are running on.
                               ''')
    ui_args.add_argument('--expose', dest="expose", 
                    help="The IP address of the server running Aqueduct.") 

    install_args = subparsers.add_parser('install', help=
                             '''Install the required library dependencies for
                             an Aqueduct connector to a third-party system.''')
    install_args.add_argument('system', nargs=1, help="Supported integrations: postgres, mysql, sqlserver, s3, snowflake, bigquery.")

    apikey_args = subparsers.add_parser('apikey', help="Display your Aqueduct API key.")
    clear_args = subparsers.add_parser('clear', help="Erase your Aqueduct installation.")

    args = parser.parse_args()

    if not os.path.isdir(base_directory):
        os.mkdir(base_directory)
    
    if args.command == "server":
        server()
    elif args.command == "install":
        install(args.system[0]) # argparse makes this an array so only pass in value [0].
    elif args.command == "ui":
        ui(args.expose)
    elif args.command == "apikey":
        apiKey()
    elif args.command == "clear":
        clear()
    else:
        print("Unsupported command:", command)
        sys.exit(1)