# 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
from src.utils.console import eprint


def assume_role_direct(role_arn: str, external_id: str, aws_session) -> dict:
    """Assume an IAM role directly without connection lookup.

    Args:
        role_arn: The ARN of the IAM role to assume
        external_id: The external ID required for the role assumption
        aws_session: AWS session object with configured credentials

    Returns:
        dict: AWS credentials (AccessKeyId, SecretAccessKey, SessionToken)
    """
    eprint(f"Assuming role directly: {role_arn}")
    try:
        sts_connection = aws_session.session.client("sts")
        response = sts_connection.assume_role(
            RoleArn=role_arn,
            RoleSessionName="cz-ops-direct",
            ExternalId=external_id,
        )
    except botocore.exceptions.ClientError as err:
        eprint("\nClient Error: Unable to assume role")
        eprint(err)
        sys.exit(1)

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


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()

    # Check if using direct role ARN mode
    role_arn = arguments.get("--role-arn")
    external_id = arguments.get("--external-id")

    if role_arn:
        # Direct role ARN mode - external_id is required
        if not external_id:
            eprint("Error: --external-id is required when using --role-arn")
            eprint("Usage: czops shell --role-arn <ARN> --external-id <ID> [--profile <PROFILE>] [--eval]")
            sys.exit(1)
        assumed_credentials = assume_role_direct(role_arn, external_id, aws_session)
    elif arguments.get("<CONNECTION-ID>"):
        # Connection ID mode - lookup connection details via API
        connection = get_connection(arguments["<CONNECTION-ID>"]).model_dump()
        assumed_credentials = get_aws_credentials_for_account(connection, aws_session, arguments.get("--role-type"))
    else:
        # Neither mode specified
        eprint("Error: Either CONNECTION-ID or --role-arn with --external-id is required")
        eprint("Usage:")
        eprint("  czops shell <CONNECTION-ID> [--profile <PROFILE>] [--eval]")
        eprint("  czops shell --role-arn <ARN> --external-id <ID> [--profile <PROFILE>] [--eval]")
        sys.exit(1)

    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")
        user_shell = os.environ.get("SHELL", "/bin/sh")
        pty.spawn(user_shell)
        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)
    account_name = account.get("connection_name", "Unknown")
    account_id = account.get("connection_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):
    provider = account.get("provider", "")
    if provider.lower() not in ("aws", "snowflake"):
        eprint("I can't find a role to assume for this account, is the account record correct?")
        sys.exit(1)

    types = account.get("types") or {}

    # Determine which type to pull credentials from
    if role_type and role_type.upper() in types:
        type_key = role_type.upper()
    elif "COST" in types:
        type_key = "COST"
    elif "RESOURCE" in types:
        type_key = "RESOURCE"
    else:
        eprint(f"No connection types found for account {account.get('connection_id', 'Unknown')}")
        sys.exit(1)

    connection_method = (types.get(type_key) or {}).get("connection_method") or {}
    role_arn = connection_method.get("role_arn")
    external_id = connection_method.get("external_id")

    if not role_arn or not external_id:
        eprint(f"Missing role_arn or external_id in {type_key} connection_method")
        eprint(f"Available types: {list(types.keys())}")
        sys.exit(1)

    return external_id, role_arn
