#!python

import json
import sys
import os
import configparser
from difflib import unified_diff
import tempfile
import yaml
import copy


def main():
    """docstring for main"""
    import argparse
    parser = argparse.ArgumentParser(
        description='Generates iterm2 dynamic profiles'
    )

    parser.add_argument('-d', '--diff',
                        action='store_true',
                        help='Print a diff to the current profiles')
    parser.add_argument('-n', '--dry-run',
                        action='store_true',
                        help='Dry run mode')
    parser.add_argument('-v', '--verbose',
                        action='count',
                        default=0,
                        help='Increase verbosity')

    args = parser.parse_args()

    aws_profiles = create_aws_profiles()
    k8s_profiles = create_k8s_profiles()

    profiles = {
        "Profiles": aws_profiles + k8s_profiles,
    }

    dest = os.path.expanduser(
        "~/Library/Application Support/iTerm2/DynamicProfiles/aws-profiles.json"
    )
    if args.diff:
        with open(dest) as input_:
            current = json.load(input_)
        sys.stdout.writelines(
            unified_diff(
                json.dumps(current, indent=4), json.dumps(profiles, indent=4)
            )
        )
    elif args.dry_run:
        print(json.dumps(profiles, indent=4))
    else:
        with open(dest, "w") as out:
            out.write(json.dumps(profiles, indent=4))


def create_k8s_profiles():
    """docstring for create_k8s_profiles"""
    user = os.getenv("USER")

    with open(os.path.expanduser("~/.kube/config")) as file:
        config = yaml.full_load(file)

    clusters = [x['name'] for x in config['clusters']]

    profiles = []
    for cluster in clusters:
        this = carve_k8s_cluster(config, cluster)
        cfg = os.path.expanduser(
            f"~/.kube/config.{cluster.replace('/', '__')}.yml"
        )
        aws_profile = None
        try:
            env = this['users'][0]['user']['exec']['env'][0]
            if env['name'] != 'AWS_PROFILE':
                pass
            aws_profile = env['value']
        except KeyError:
            pass
        except TypeError:
            pass
        except IndexError:
            pass

        with open(cfg, "w") as out:
            yaml.dump(this, out)
        new = create_profile(
            f'k8s-{cluster}',
            cmd=f"/usr/bin/env KUBECONFIG={cfg} AWS_PROFILE={aws_profile} /usr/bin/login -fp {user}",
            change_title=False,
            tags=['k8s'],
        )
        new["Background Color"] = {
            "Red Component": 0,
            "Color Space": "sRGB",
            "Blue Component": 0.38311767578125,
            "Alpha Component": 1,
            "Green Component": 0
        }
        profiles += [new]
    return profiles


def carve_k8s_cluster(config, cluster):
    """docstring for carve_k8s_cluster"""

    config_cluster = [x for x in config['clusters'] if x['name'] == cluster]
    config_context = [x for x in config['contexts'] if x['name'] == cluster]
    config_user = [x for x in config['users'] if x['name'] == cluster]

    new = copy.deepcopy(config)
    new['clusters'] = config_cluster
    new['contexts'] = config_context
    new['users'] = config_user
    new['current-context'] = cluster

    return new


def create_aws_profiles():
    aws_profiles = get_aws_profiles()

    profiles = []
    for prof, account, role, loggin_for in aws_profiles:
        profiles += [mkprofile(prof, account, role, loggin_for)]
    return profiles


def create_profile(name, cmd=None, change_title=False, tags=None):
    if tags is None:
        tags = []

    ret = {
        "Name": name,
        "Guid": name,
        "Unlimited Scrollback": True,
        "Title Components": 32,
        "Custom Window Title": name,
        "Allow Title Setting": change_title,
        "Tags": tags,
        "Badge Text": name,
    }

    if cmd is not None:
        ret["Command"] = cmd
        ret["Custom Command"] = "Yes"
    return ret


def mkprofile(aws_profile, account=None, role=None, loggin_for=None):
    user = os.getenv("USER")
    ret = create_profile(
        aws_profile,
        cmd=f"/usr/bin/env AWS_PROFILE={aws_profile} /usr/bin/login -fp {user}",
        change_title=False,
    )

    if account is not None:
        ret['Tags'] += [account]
    if role is not None:
        ret['Tags'] += [role]
    if loggin_for is not None:
        ret['Tags'] += [loggin_for]

    if 'prod' in aws_profile:
        ret["Background Color"] = {
            "Red Component": 0.217376708984375,
            "Color Space": "sRGB",
            "Blue Component": 0,
            "Alpha Component": 1,
            "Green Component": 0
        }
    return ret


def get_aws_profiles():
    config = configparser.ConfigParser()
    config.read(os.path.expanduser('~/.aws/config'))
    ret = []
    loggin_profiles = {}
    for x in config.sections():
        account_number = None
        role = None
        name = x.split()[1]
        if 'role_arn' in config[x]:
            account_number = config[x]['role_arn'].split(':')[4]
            role = config[x]['role_arn'].split(':')[5]
        if 'source_profile' in config[x]:
            loggin_profiles[config[x]['source_profile']] = name
        ret += [[name, account_number, role]]

    for x in ret:
        log = None
        if x[0] in loggin_profiles.keys():
            log = f'log-{loggin_profiles[x[0]]}'
        x += [log]
    return ret


if __name__ == '__main__':
    main()
