#!/usr/bin/env python
import ecssh
import boto3
import yaml
import sys
import os
import subprocess
import shlex


def loadConfig(CLUSTER):
    with open((os.getenv("HOME") + "/.config/ecssh.yml"), 'r') as stream:
        try:
            RAW_CONFIG = yaml.safe_load(stream)
        except yaml.YAMLError as exc:
            print(exc)
            sys.exit(1)

    PROFILE_CONFIG = RAW_CONFIG.get(AWS_PROFILE)
    CONFIG = PROFILE_CONFIG.get(CLUSTER)
    if CONFIG.get("ssh_key"):
        CONFIG["ssh_key"] = "-i " + CONFIG.get("ssh_key")
    return CONFIG


AWS_PROFILE = os.getenv("AWS_PROFILE", "default")
session = boto3.Session(profile_name=AWS_PROFILE)
ECS = session.client('ecs')
EC2 = session.client('ec2')


if __name__ == "__main__":
    if len(sys.argv) == 2:
        CLUSTER = sys.argv[1]
        CONFIG = loadConfig(CLUSTER)
        ecssh.ecssh(CONFIG=CONFIG, ECS=ECS,
                    EC2=EC2, CLUSTER=CLUSTER, N=None)
    elif len(sys.argv) == 3:
        CLUSTER = sys.argv[1]
        N = int(sys.argv[2])
        CONFIG = loadConfig(CLUSTER)
        subprocess.check_call(shlex.split(ecssh.ecssh(
            CONFIG=CONFIG, ECS=ECS, EC2=EC2, CLUSTER=CLUSTER, N=N)))
    else:
        print(""" USAGE:
        
        AWS_PROFILE=<ENV> ecssh <CLUSTER NANE>
        AWS_PROFILE=<ENV> ecssh <CLUSTER NANE> <N>
        """)
