#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import subprocess
from typing import List

LIVE_PROJECT_NAME = "my-project-live"
STG_PROJECT_NAME = "my-project-stg"

SRC_DB_NAME = "project-db"
SRC_DB_USER_NAME = "project-db-username"

DST_DB_NAME = "project-db"
DST_DB_USER_NAME = "project-db-username"
DST_DB_PASSWORD = ""
DST_DB_PORT = "5432"

DUMP_FILE_NAME = "/tmp/database-dump.sql"
DUMP_ALTERED_FILE_NAME = "/tmp/database-dump-altered.sql"

regions = {
    "eu": "europe-west1",
    "us": "us-east1",
}

clusters = {
    "eu": "gke_my-project-live_europe-west1_01",
    "us": "gke_my-project-live_us-east1_01",
}

countries_live: List[Country] = [
    Country(name="USA", code="us", region="us", db_instance="live-us"),
    Country(name="Germany", code="de", region="eu", db_instance="live-de"),
]

countries_stg: List[Country] = [
    Country(name="USA", code="us", region="us", db_instance="stg-us", project=STG_PROJECT_NAME),
]

countries_local: List[Country] = [
    Country(name="Local", code="lo", region="lo", db_instance="local", project=""),
]

class Country:
    name: str = ""
    code: str = ""
    region: str = ""
    db_instance: str = ""
    db_instance_replica: str = ""
    project_name: str = LIVE_PROJECT_NAME

    def __init__(self, name: str = "", code: str = "", region: str = "", db_instance: str = "", db_instance_replica: str = "", canary: bool = False, project_name: str = LIVE_PROJECT_NAME):
        self.name = name
        self.code = code
        self.region = region
        self.db_instance = db_instance
        self.db_instance_replica = db_instance_replica
        if self.db_instance_replica == "":
            self.db_instance_replica = self.db_instance+"-replica"
        self.project_name = project_name
        pass

    def get_connection(self) -> str:
        return "{}:{}:{}".format(self.project_name, regions[self.region], self.db_instance)

    def get_replica_connection(self) -> str:
        return "{}:{}:{}".format(self.project_name, regions[self.region], self.db_instance_replica)

    def get_name_upper(self) -> str:
        return self.name.upper()

    def get_code_upper(self) -> str:
        return self.code.upper()


def main():
    args = parse_arguments()
    action = args.action
    env = args.env
    country = args.country
    rw = args.rw

    if action == "proxy_database":
        found_country = find_country(country, env)
        if found_country == None:
            print("No such country: {}".format(country))
            exit(1)
        if rw:
            print("Caution: read-write mode enabled!")
        proxy_database(found_country, rw)
    elif action == "select_cluster":
        found_country = find_country(country, env)
        if found_country == None:
            print("No such country: {}".format(country))
            exit(1)
        select_cluster(found_country)
    elif action == "dump_database":
        dump_database()

    exit(0)


def select_cluster(country: Country):
    kubectl_cmd = get_absolute_path("kubectl")
    execute("{} config use-context {}".format(kubectl_cmd, clusters[country.region]))


def proxy_database(country: Country, rw: bool):
    connection = country.get_connection()
    if rw:
        connection = country.get_replica_connection()

    cloud_sql_proxy_cmd = get_absolute_path("cloud_sql_proxy")
    execute("{} -instances={}=tcp:54321".format(cloud_sql_proxy_cmd, connection))


def dump_database() -> None:
    print("This action dumps the remote database to the local Postgres instance. Make sure the SQL proxy is running.")

    src_db_password = getpass.getpass("Source database password: ")
    if src_db_password == "":
        print("The password may not be empty")
        exit(1)

    src_connection_string = "postgres://{}:{}@127.0.0.1:54321/{}".format(SRC_DB_USER_NAME, src_db_password, SRC_DB_NAME)
    dst_connection_string = "postgres://{}:{}@127.0.0.1:{}".format(DST_DB_USER_NAME, DST_DB_PASSWORD, DST_DB_PORT)
    pg_connection_string = "{}/postgres".format(dst_connection_string)

    print("Running this action will overwrite your local database. Type \"yes\" to proceed.")
    user_confirmation = input()

    if user_confirmation != "yes":
        print("Operation cancelled")
        exit(0)

    print("Dumping the remote database...")
    execute_no_echo("pg_dump {} > {}".format(src_connection_string, DUMP_FILE_NAME))
    execute("sed -e \"s/OWNER TO {}/OWNER TO {}/g\" {} > {}".format(SRC_DB_USER_NAME, DST_DB_USER_NAME, DUMP_FILE_NAME, DUMP_ALTERED_FILE_NAME))
    execute("psql \"{}\" -c \"drop database if exists \\\"{}\\\"\"".format(pg_connection_string, DST_DB_NAME))
    execute("psql \"{}\" template1 -c \"create database \\\"{}\\\" with owner {}\"".format(pg_connection_string, DST_DB_NAME, DST_DB_USER_NAME))
    print("Importing the database locally...")
    execute("psql \"{}/{}\" < {}".format(dst_connection_string, DST_DB_NAME, DUMP_ALTERED_FILE_NAME))
    execute("rm {}".format(DUMP_FILE_NAME))
    execute("rm {}".format(DUMP_ALTERED_FILE_NAME))
    print("Done!")


def execute(command):
    print("Executing: {}".format(command))
    try:
        subprocess.run(command, shell=False)
    except KeyboardInterrupt:
        print("Ctrl+C pressed. Exiting gracefully.")


def execute_no_echo(command):
    try:
        subprocess.run(command, shell=False)
    except KeyboardInterrupt:
        print("Ctrl+C pressed. Exiting gracefully.")


def get_absolute_path(cmd):
    output = subprocess.run("which {}".format(cmd), capture_output=True, shell=True)
    return output.stdout.decode()[:-1]


def find_country(country: str, env: str) -> Country:
    country_list = countries_stg
    if env = "live":
        country_list = countries_live
    elif env = "lo":
        country_list = countries_local

    country_uppercased = country.upper()
    for country_element in countries:
        if (country_element.get_name_upper() == country_uppercased or country_element.get_code_upper() == country_uppercased):
            return country_element

    return None


def parse_arguments():
    parser = argparse.ArgumentParser(description="A helper tool for faster devops")

    parser.add_argument("country", type=str, help="The country name or code")
    parser.add_argument("env", choices=["live", "stg", "lo"], help="Environment")
    parser.add_argument("action", choices=["proxy_database", "select_cluster", "dump_database"], help="What to do")
    parser.add_argument("--rw", action="store_true", help="If set to TRUE, the connection will be writable")

    return parser.parse_args()


if __name__ == "__main__":
    main()