# Copyright (c) 2020 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com
#

import os
import pty
import sys

import botocore
from rich.pretty import pprint

from src.api.v4_connections import get_connection_for_shell
from src.utils.console import eprint


def create_shell(arguments, aws_session):
    if arguments.get("--reset"):
        print("unset AWS_ACCESS_KEY_ID")
        print("unset AWS_SECRET_ACCESS_KEY")
        print("unset AWS_SESSION_TOKEN")
        eprint("Your AWS Session has been reset")
        sys.exit()

    # Validate required parameters for non-reset operations
    if not arguments.get("<ID>"):
        eprint("Error: --connection-id is required when not using --reset")
        sys.exit(1)

    connection = get_connection_for_shell(arguments["<ID>"])
    assumed_credentials = get_aws_credentials_for_account(connection, aws_session, arguments.get("--role-type"))

    if arguments["--eval"]:
        print(f"export AWS_ACCESS_KEY_ID={assumed_credentials['AccessKeyId']}")
        print(f"export AWS_SECRET_ACCESS_KEY={assumed_credentials['SecretAccessKey']}")
        print(f"export AWS_SESSION_TOKEN={assumed_credentials['SessionToken']}")
    else:
        os.environ["AWS_ACCESS_KEY_ID"] = assumed_credentials["AccessKeyId"]
        os.environ["AWS_SECRET_ACCESS_KEY"] = assumed_credentials["SecretAccessKey"]
        os.environ["AWS_SESSION_TOKEN"] = assumed_credentials["SessionToken"]
        os.environ["AWS_DEFAULT_REGION"] = "us-east-1"
        os.environ["AWS_DEFAULT_OUTPUT"] = "json"
        os.environ["TERM"] = "xterm-256color"
        os.environ["PS1"] = "$ "
        print("'crtl-d <enter>' to exit")
        pty.spawn("/bin/sh")
        print("Shell closed")


def get_aws_credentials_for_account(account, aws_session, role_type=None):
    external_id, role_arn = get_access_role_for_account(account, role_type)
    # V4 API uses connection_name and connection_id, legacy uses cz_account_name and cloud_account_id
    account_name = account.get("connection_name") or account.get("cz_account_name", "Unknown")
    account_id = account.get("connection_id") or account.get("cloud_account_id", "Unknown")
    eprint(f"Getting credentials for {account_name} ({account_id})")
    eprint(f"Role: {role_arn}")
    try:
        sts_connection = aws_session.session.client("sts")
        response = sts_connection.assume_role(RoleArn=role_arn, RoleSessionName="cz-ops", ExternalId=external_id)

    except botocore.exceptions.ClientError as err:
        eprint("\nClient Error: Unable to assume role and connect to account")
        eprint(err)
        eprint("Verify you have selected a role that can assume roles from the CZ Reactor account")
        sys.exit(1)
    except Exception as err:
        eprint("\nUNKNOWN ERROR: Unable to assume role and connect to account")
        eprint(err)
        pprint(aws_session)
        sys.exit(1)

    assumed_credentials = response.get("Credentials")
    if not assumed_credentials:
        eprint(" - FAILED to assume role, check your permissions")
        eprint(response)
        sys.exit()
    return assumed_credentials


def get_access_role_for_account(account, role_type=None):
    if role_type is None:
        role_type = "---"

    # V4 API stores provider as "provider", legacy as "cloud_provider"
    provider = account.get("provider") or account.get("cloud_provider", "")

    # V4 API stores role details in connection_details, legacy at top level
    connection_details = account.get("connection_details", {})

    if provider.lower() == "aws":
        if role_type == "resource":
            role_arn = connection_details.get("resource_owner_role_arn") or account.get("resource_owner_role_arn")
        elif role_type == "payer":
            role_arn = connection_details.get("role_arn") or account.get("master_payer_role_arn")
        elif account.get("role_type"):
            role_arn = account[role_type]
        elif connection_details.get("is_cost_connection") or account.get("is_cost_connection", False):
            role_arn = connection_details.get("role_arn") or account.get("master_payer_role_arn")
        else:
            role_arn = connection_details.get("role_arn") or account.get("resource_owner_role_arn")
    elif provider.lower() == "snowflake":
        role_arn = connection_details.get("role_arn") or account.get("master_payer_role_arn")
    else:
        eprint("I can't find a role to assume for this account, is the account record correct?")
        sys.exit()

    external_id = connection_details.get("external_id") or account.get("external_id")
    return external_id, role_arn
