package setup

import (
	"fmt"

	"github.com/fatih/color"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"

	"crypto/md5"
	"database/sql"
	"os"
	"os/exec"
	"strconv"

	"code.justin.tv/d8a/buddy/cmd/buddy-cli/data"
	"code.justin.tv/d8a/buddy/lib/alerts"
	"code.justin.tv/d8a/buddy/lib/clusters"
	"code.justin.tv/d8a/buddy/lib/clusters/clusterdb"
	"code.justin.tv/d8a/buddy/lib/config"
	"code.justin.tv/d8a/buddy/lib/git"
	"code.justin.tv/d8a/buddy/lib/sandstorm"
	"code.justin.tv/d8a/buddy/lib/store"
	"code.justin.tv/systems/sandstorm/manager"
	"code.justin.tv/twitch/cli/ghutil"

	"io/ioutil"
	"strings"
	"time"

	"code.justin.tv/d8a/buddy/lib/zenyatta"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"github.com/aws/aws-sdk-go/service/rds"
	"github.com/aws/aws-sdk-go/service/sns"
)

var setupCmd = &cobra.Command{
	Use:   "setup",
	Short: "Execute an iceman command against a target cluster",
	Long:  "",
	RunE: func(cmd *cobra.Command, args []string) error {
		configFile := data.RootData.ConfigFile()
		client, err := ghutil.ClientFromToken(configFile.GithubKey)
		if err != nil {
			return err
		}

		fmt.Println("Locating buddy's datastore...")
		rdsInstance, err := data.RootData.LocateDataStore(true)
		if err != nil {
			return err
		}

		fmt.Println("Running iceman against datastore...")
		err = store.MigrateStore(rdsInstance, data.RootData.SandstormClient())
		if err != nil {
			return err
		}
		fmt.Println("Iceman complete.")

		db, err := store.OpenDbConn(rdsInstance, data.RootData.SandstormClient())
		if err != nil {
			return err
		}
		fmt.Println("Connected to datastore.")

		fmt.Println("Setting up redis for airflow.")
		err = setupZenyattaRedis(configFile.SandstormTeam)
		if err != nil {
			return errors.Wrap(err, "airflow failed to init")
		}

		fmt.Println("Writing airflow.cfg")
		if err = reWriteAirflowConfig(); err != nil {
			return errors.Wrap(err, "could not rewrite airflow.cfg")
		}

		fmt.Println("Pushing zenyatta config to S3")
		if err = pushZenyattaConfig("etc/zenyatta"); err != nil {
			return errors.Wrap(err, "could not push zenyatta config")
		}

		fmt.Println("Initializing airflow DB")
		err = initAirflow(db)
		if err != nil {
			return errors.Wrap(err, "airflow failed to init")
		}

		fmt.Println("Setting up worker nodes for airflow.")
		err = zenyatta.StartZenyattaWorker(configFile.SandstormTeam, data.RootData.Region(), configFile.Environment, "etc/zenyatta")
		if err != nil {
			return errors.Wrap(err, "airflow failed to spin up worker node")
		}

		fmt.Println("Creating KeyPair for EMR cluster for airflow if not exists.")
		err = createKeyPairEMR(configFile.SandstormTeam, configFile.Environment, data.RootData.SandstormClient())
		if err != nil {
			return errors.Wrap(err, "buddy failed to create key pair for airflow emr cluster")
		}

		fmt.Println("Setting up EMR cluster for airflow.")
		err = setupZenyattaEMR(configFile.SandstormTeam, configFile.Environment)
		if err != nil {
			return errors.Wrap(err, "airflow failed to spin up EMR cluster")
		}

		configFile.Cluster, err = store.GetClusters(db)
		if err != nil {
			return err
		}

		fmt.Println("Retrieving config cluster data & replicas...")
		instances, rdsClusters, err := clusters.FetchMeta(configFile.Cluster, data.RootData.RdsClient())
		if err != nil {
			return err
		}
		fmt.Println("Data retrieved")

		fmt.Println("Ensuring instance cloudwatch alerts are live...")
		instancePriorities := make(map[string]alerts.AlarmPriority)
		clusterPriorities := make(map[string]alerts.AlarmPriority)
		for _, cluster := range configFile.Cluster {
			clusterPriorities[cluster.RootIdentifier] = cluster.AlarmPriority
		}
		for instanceName, instance := range instances {
			instancePriorities[instanceName] = getAlarmPriority(instance, clusterPriorities, instances, rdsClusters)
		}

		cwClient := cloudwatch.New(data.RootData.Session())
		snsClient := sns.New(data.RootData.Session())
		err = alerts.ConfigureAlarms(snsClient, cwClient, configFile.SandstormTeam, instancePriorities, instances)
		if err != nil {
			color.Red("%v", err)
		} else {
			fmt.Println("Alerts live")
		}

		//We're going to simultaneously set up all clusters because there are some long waits involved
		err = clusters.ProcessInParallel(configFile.Cluster,
			func(cluster *config.Cluster) (interface{}, error) {
				err = git.EnsureRepository(configFile, cluster, client)
				return cluster.Repository, err
			},
			func(result interface{}, err error) error {
				if err != nil {
					color.Red("%v", err)
				} else {
					repoText, ok := result.(string)
					if ok {
						color.Green("Successfully configured repository %v for iceman.", repoText)
					} else {
						color.Yellow("%v", result)
					}
				}

				return nil
			})
		if err != nil {
			color.Red("%v", err)
		}

		fmt.Println("Syncing sandstorm to users")
		clusterMap := make(map[string]*config.Cluster)
		for _, cluster := range configFile.Cluster {
			clusterMap[cluster.Name] = cluster
		}
		changed, err := sandstorm.RemoveDeadClusterTemplates(clusterMap)
		if err != nil {
			color.Red("%v", err)
		}

		err = clusters.ProcessInParallel(configFile.Cluster,
			func(cluster *config.Cluster) (interface{}, error) {
				clusterDb, err := clusterdb.OpenDbConn(cluster, data.RootData.SandstormClient(), cluster.SuperUser)
				if err != nil {
					return false, errors.Wrap(err, fmt.Sprintf("could not open connection to cluster %s", cluster.Name))
				}
				superPassword, err := data.RootData.SandstormClient().GetUserPassword(cluster, cluster.SuperUser)

				if err != nil {
					return false, errors.Wrap(err, fmt.Sprintf("could not get password for %s on %s", cluster.SuperUser, cluster.Name))
				}

				err = addAirflowConnection(cluster.Name, cluster.Driver, cluster.Host, cluster.SuperUser, superPassword, cluster.Database, cluster.Port)

				if err != nil {
					return false, errors.Wrap(err, fmt.Sprintf("could not add airflow connection for %s", cluster.Name))
				}

				users, err := clusterDb.ClusterUsers()

				if err != nil {
					return false, errors.Wrap(err, fmt.Sprintf("could not retrieve users from cluster %s", cluster.Name))
				}

				return sandstorm.SyncClusterTemplates(data.RootData.SandstormClient(), cluster, users)
			},
			func(result interface{}, err error) error {
				boolResult, ok := result.(bool)
				if ok {
					changed = changed || boolResult
				}
				return err
			})
		if err != nil {
			color.Red("%v", err)
		}

		fmt.Println("Setting up db schema and partition for airflow.")
		err = setupConnectionSchema()
		if err != nil {
			return errors.Wrap(err, "airflow failed to set up")
		}

		passwordsChanged, err := sandstorm.WritePasswordFiles(configFile.Cluster)
		if err != nil {
			color.Red("%v", err)
		}
		changed = changed || passwordsChanged

		if changed {
			err = sandstorm.Reload()
			if err != nil {
				color.Red("%v", err)
			}
		}

		if err != nil {
			fmt.Println("Sandstorm synched")
		}

		return nil
	},
}

func getAlarmPriority(instance *rds.DBInstance, clusterPriorities map[string]alerts.AlarmPriority, instances map[string]*rds.DBInstance, clusters map[string]*rds.DBCluster) alerts.AlarmPriority {
	if instance.ReadReplicaDBClusterIdentifiers != nil {
		cluster, ok := clusters[*instance.DBClusterIdentifier]
		if !ok {
			return alerts.NoAlarms
		}

		return getAlarmPriorityForCluster(cluster, clusterPriorities, instances, clusters)
	}

	if instance.ReadReplicaSourceDBInstanceIdentifier != nil {
		instanceMaster, ok := instances[*instance.ReadReplicaSourceDBInstanceIdentifier]
		if !ok {
			return alerts.NoAlarms
		}

		return getAlarmPriority(instanceMaster, clusterPriorities, instances, clusters)
	}

	priority, _ := clusterPriorities[*instance.DBInstanceIdentifier]
	return priority
}

func getAlarmPriorityForCluster(cluster *rds.DBCluster, clusterPriorities map[string]alerts.AlarmPriority, instances map[string]*rds.DBInstance, clusters map[string]*rds.DBCluster) alerts.AlarmPriority {
	if cluster.ReplicationSourceIdentifier != nil {
		masterCluster, ok := clusters[*cluster.ReplicationSourceIdentifier]
		if !ok {
			return alerts.NoAlarms
		}

		return getAlarmPriorityForCluster(masterCluster, clusterPriorities, instances, clusters)
	}

	priority, _ := clusterPriorities[*cluster.DBClusterIdentifier]
	return priority
}

func reWriteAirflowConfig() error {
	configPath := "/etc/sandstorm-agent/templates.d/zenyatta-airflow-config"
	fileAsBytes, err := ioutil.ReadFile(configPath)
	if err != nil {
		return errors.Wrap(err, "could not open "+configPath+" for reading")
	}
	inputLines := strings.Split(string(fileAsBytes), "\n")

	var outputLines []string
	for _, line := range inputLines {
		if strings.HasPrefix(line, "sql_alchemy_conn") {
			connString, err := AirflowSqlalchemyConnection()
			if err != nil {
				return errors.Wrap(err, "couldn't get sqlalchemy connection string")
			}
			line = fmt.Sprintf("sql_alchemy_conn = %s", connString)
		} else if strings.HasPrefix(line, "broker_url") {
			connString, err := AirflowRedisConnection()
			if err != nil {
				return errors.Wrap(err, "couldn't get redis connection string")
			}
			line = fmt.Sprintf("broker_url = %s", connString)
		} else if strings.HasPrefix(line, "celery_result_backend") {
			connString, err := AirflowRedisConnection()
			if err != nil {
				return errors.Wrap(err, "couldn't get redis connection string")
			}
			line = fmt.Sprintf("celery_result_backend = %s", connString)
		} else if strings.HasPrefix(line, "executor") {
			line = "executor = CeleryExecutor"
		}
		outputLines = append(outputLines, line)
	}

	err = sandstorm.Reload()
	if err != nil {
		return errors.Wrap(err, "couldn't reload sandstorm, and the airflow config template is not in a good state")
	}
	err = ioutil.WriteFile(configPath, []byte(strings.Join(outputLines, "\n")), 0644)
	if err != nil {
		return errors.Wrap(err, "couldn't write "+configPath)
	}
	sqlString, err := AirflowSqlalchemyConnection()
	if err != nil {
		return errors.Wrap(err, "couldn't get sqlalchemy connection string")
	}

	fileNotWritten := true
	airflowPath := "/etc/zenyatta/airflow.cfg"
	attempts := 10
	for fileNotWritten && attempts > 0 {
		attempts--
		theBytes, err := ioutil.ReadFile(airflowPath)
		if err != nil {
			return errors.Wrap(err, "couldn't read airflow.cfg")
		}
		theString := string(theBytes)
		if strings.Contains(theString, sqlString) {
			fileNotWritten = false
		} else {
			time.Sleep(6000 * time.Millisecond)
		}
	}

	if fileNotWritten == true {
		return errors.New("failed to detect changes in " + airflowPath + "; check sandstorm")
	}
	return nil
}

func initAirflow(db *sql.DB) error {
	fmt.Println("potentially creating airflow user")
	// get airflow password from sandstorm
	if _, err := db.Exec(`
	DO
	$body$
	BEGIN
   		IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'airflow') THEN
      			CREATE ROLE airflow;
   		END IF;
	END
	$body$;
	`); err != nil {
		return errors.Wrap(err, "blew up at create user")
	}

	configFile := data.RootData.ConfigFile()
	secretPath := fmt.Sprintf("%s/rds-buddy-store/%s/airflow_password", configFile.SandstormTeam, configFile.Environment)
	password, err := data.RootData.SandstormClient().Manager().Get(secretPath)
	if err != nil {
		return errors.Wrap(err, "couldn't get airlfow password from sandstorm")
	}
	encryptedPass := md5.Sum([]byte(string(password.Plaintext) + "airflow"))
	query := fmt.Sprintf("ALTER USER airflow LOGIN ENCRYPTED PASSWORD 'md5%x'", encryptedPass)
	_, err = db.Exec(query)
	if err != nil {
		return errors.Wrap(err, "blew up at alter user")
	}
	result, err := db.Query("SELECT 1 FROM pg_database WHERE datname = 'airflow'")
	if err != nil {
		return errors.Wrap(err, "unable to SELECT 1 from pg_database")
	}
	tableResult := 0
	for result.Next() {
		err = result.Scan(&tableResult)
		if err != nil {
			return errors.Wrap(err, "failed to scan result of query")
		}
	}
	if tableResult != 1 {
		_, err := db.Exec("CREATE DATABASE airflow")
		if err != nil {
			return errors.Wrap(err, "blew up at create database")
		}
	}
	if _, err = db.Exec("GRANT ALL PRIVILEGES ON DATABASE airflow to airflow"); err != nil {
		return errors.Wrap(err, "couldn't grant privs")
	}

	env, err := zenyatta.ZenyattaEnvironment()
	if err != nil {
		return err
	}

	airflowInitDb := exec.Command("/opt/virtualenvs/airflow/bin/airflow", "initdb")
	airflowInitDb.Env = env
	airflowInitDb.Stdout = os.Stdout
	airflowInitDb.Stderr = os.Stderr
	airflowInitDb.Stdin = os.Stdin
	err = airflowInitDb.Run()
	if err != nil {
		return errors.Wrap(err, "failed to run airflow initdb")
	}

	return nil
}

func addAirflowConnection(connId, connType, host, user, password, schema string, port int) error {

	cmd := exec.Command("/opt/virtualenvs/airflow/bin/python",
		"/opt/twitch/zenyatta/current/scripts/add_etl_dag.py",
		"--rds_id", connId,
		"--type", connType,
		"--host", host,
		"--user", user,
		"--password", password,
		"--schema", schema,
		"--port", strconv.Itoa(port))
	env, err := zenyatta.ZenyattaEnvironment()
	if err != nil {
		return err
	}
	cmd.Env = env
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	err = cmd.Run()
	fmt.Println(fmt.Sprintf("adding %s to zenyatta", connId))
	if err != nil {
		return errors.Wrap(err, "failed to run add_etl_dag script, this is likely because the connection exists in the airflow db already")
	}

	return nil
}

func setupZenyattaRedis(teamName string) error {

	cmd := exec.Command("/opt/virtualenvs/airflow/bin/python",
		"/opt/twitch/zenyatta/current/scripts/setup_redis.py",
		"--name", teamName,
		"--region", data.RootData.Region())
	env, err := zenyatta.ZenyattaEnvironment()
	if err != nil {
		return err
	}

	cmd.Env = env
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	err = cmd.Run()
	if err != nil {
		return errors.Wrap(err, "failed to run setup_redis script")
	}

	return nil
}

func setupZenyattaEMR(teamName string, environment string) error {

	cmd := exec.Command("/opt/virtualenvs/airflow/bin/python",
		"/opt/twitch/zenyatta/current/scripts/setup_emr.py",
		"--name", teamName,
		"--region", data.RootData.Region(),
		"--env", environment)
	env, err := zenyatta.ZenyattaEnvironment()
	if err != nil {
		return err
	}

	cmd.Env = env
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	err = cmd.Run()
	if err != nil {
		return errors.Wrap(err, "failed to run setup_emr.py script")
	}

	return nil
}

func pushZenyattaConfig(key string) error {
	// defaults to copying /etc/zenyatta to the s3://bucket/key path
	cmd := exec.Command("/opt/virtualenvs/airflow/bin/python",
		"/opt/twitch/zenyatta/current/scripts/push_pull_config.py",
		"--key", key,
		"push")
	env, err := zenyatta.ZenyattaEnvironment()
	if err != nil {
		return err
	}

	cmd.Env = env
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	err = cmd.Run()
	if err != nil {
		return errors.Wrap(err, "failed to run push_pull_config.py")
	}

	return nil

}

func setupConnectionSchema() error {

	cmd := exec.Command("/opt/virtualenvs/airflow/bin/python",
		"/opt/twitch/zenyatta/current/scripts/setup_schema_partition.py")
	env, err := zenyatta.ZenyattaEnvironment()
	if err != nil {
		return err
	}

	cmd.Env = env
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	err = cmd.Run()
	if err != nil {
		return errors.Wrap(err, "failed to run setup_schema_partition.py script")
	}

	return nil
}

func createKeyPairEMR(teamName string, environment string, sandstormClient sandstorm.SandstormAPI) error {

	keyPairExists, keyContent, err := data.RootData.LocateKeyPair(fmt.Sprintf("%s-rds-buddy-store-%s-emr", teamName, environment))

	if err != nil {
		return nil
	}

	if keyPairExists {
		fmt.Printf("buddy found key pair exists\n")
	} else {
		fmt.Printf("buddy created a key pair and will store it sandstorm \n")

		err = sandstormClient.Manager().Put(&manager.Secret{
			Name:      fmt.Sprintf("%s/rds-buddy-store/%s/emr", teamName, environment),
			Plaintext: []byte(keyContent),
		})
		if err != nil {
			fmt.Println(fmt.Sprintf("failed to store key pair in sandstorm: %v", err))
			return nil
		}
	}

	return nil
}

func init() {
	RootCmd.AddCommand(setupCmd)
}
