import sys

import boto3
import requests
from requests_aws4auth import AWS4Auth
import os


def get_auth(region: str):
    service = 'es'
    credentials = boto3.Session().get_credentials()
    return AWS4Auth(
        credentials.access_key,
        credentials.secret_key,
        region,
        service,
        session_token=credentials.token
    )


def register_repo(host: str, region: str, awsauth, bucket_name: str, repository: str, backup_role_arn: str):
    headers = {"Content-Type": "application/json"}
    path = f"_snapshot/{repository}"
    url = host + path

    payload = {
        "type": "s3",
        "settings": {
            "bucket": bucket_name,
            "region": region,
            "role_arn": backup_role_arn
        }
    }

    resp = requests.put(url, auth=awsauth, json=payload, headers=headers)

    print(resp.text)

    return resp.status_code != 200


def create_snapshot(host: str, region: str, awsauth, bucket_name: str, repository_name: str, snapshot_name: str):
    path = f"_snapshot/{repository_name}/{snapshot_name}"
    url = host + path
    resp = requests.put(url, auth=awsauth)

    print(resp.text)

    return resp.status_code != 200


def restore_snapshot(host: str, awsauth, repository_name: str, snapshot_name: str):
    headers = {"Content-Type": "application/json"}
    path = f"_snapshot/{repository_name}/{snapshot_name}/_restore"
    url = host + path

    payload = {
        "indices": "-.kibana*,-.opendistro_security",
        "include_global_state": False
    }

    resp = requests.post(url, auth=awsauth, json=payload, headers=headers)

    print(resp.text)

    return resp.status_code != 200


def main(args):
    # include https:// and trailing /
    host = os.environ.get('ES_URL')
    region = os.environ.get('AWS_REGION', 'us-east-1')
    repository = os.environ.get('ES_SNAPSHOT_REPO')
    bucket_name = os.environ.get('S3_BUCKET_NAME')
    snapshot_name = os.environ.get('ES_SNAPSHOT_NAME')
    backup_role_arn = os.environ.get('ES_IAM_BACKUP_ROLE')
    awsauth = get_auth(region)

    if len(args) != 1:
        raise ValueError('Please provide exactly one argument: register, '
                         'backup or restore')

    if args[0] == 'register':
        return register_repo(host,
                             region,
                             awsauth,
                             bucket_name,
                             repository,
                             backup_role_arn)
    elif args[0] == 'restore':
        return restore_snapshot(host,
                                awsauth,
                                repository,
                                snapshot_name)
    elif args[0] == 'backup':
        return create_snapshot(host,
                               region,
                               awsauth,
                               bucket_name,
                               repository,
                               snapshot_name)
    else:
        raise ValueError('Please provide exactly one argument: register, '
                         'backup or restore')


if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))