import boto3
import json
from actions_logging.app_logging import logger
from github.env import get_required_env_var, write_github_env, exit_on_error_and_write_summary
from aws.msk.msk_info import (
    get_msk_connector_role_name,
    get_kafka_resource_arn_list,
    get_msk_environment_name
)
from aws.client import create_client_or_exit
from aws.env_info import get_env_account
import os

def get_or_create_role(role_name: str, iam_client: boto3.client) -> str:
    # Check if the role exists
    try:
        iam_client.get_role(RoleName=role_name)
        logger.info(f"Role {role_name} already exists.")
    except iam_client.exceptions.NoSuchEntityException:
        # Create the IAM role if it does not exist
        try:
            iam_client.create_role(
                RoleName=role_name,
                AssumeRolePolicyDocument=json.dumps({
                    "Version": "2012-10-17",
                    "Statement": [
                        {
                            "Sid": "",
                            "Effect": "Allow",
                            "Principal": {"Service": "kafkaconnect.amazonaws.com"},
                            "Action": "sts:AssumeRole"
                        }
                    ]
                })
            )
            logger.info(f"Role {role_name} created successfully.")
        except Exception as e:
            exit_on_error_and_write_summary(f"Failed to create role {role_name}: {e}")
    # Get role ARN
    role_details = iam_client.get_role(RoleName=role_name)
    role_arn = role_details['Role']['Arn']
    return role_arn

# Check if inline policies already exist before attaching them
def ensure_inline_policy(role_name, policy_name, policy_document, iam_client: boto3.client):
    try:
        existing_policies = iam_client.list_role_policies(RoleName=role_name)['PolicyNames']
        if policy_name in existing_policies:
            logger.info(f"Policy {policy_name} already attached to {role_name}.")
        else:
            iam_client.put_role_policy(
                RoleName=role_name,
                PolicyName=policy_name,
                PolicyDocument=json.dumps(policy_document)
            )
            logger.info(f"Policy {policy_name} attached to {role_name}.")
    except Exception as e:
        logger.error(f"Error attaching policy {policy_name} to {role_name}: {e}")

# Check if managed policy is already attached
def ensure_managed_policy(role_name, policy_arn, iam_client: boto3.client):
    try:
        attached_policies = iam_client.list_attached_role_policies(RoleName=role_name)['AttachedPolicies']
        if any(policy['PolicyArn'] == policy_arn for policy in attached_policies):
            logger.info(f"Managed policy {policy_arn} already attached to {role_name}.")
        else:
            iam_client.attach_role_policy(RoleName=role_name, PolicyArn=policy_arn)
            logger.info(f"Managed policy {policy_arn} attached to {role_name}.")
    except Exception as e:
        logger.error(f"Error attaching managed policy {policy_arn} to {role_name}: {e}")

def get_msk_role_policies(env_name: str, aws_msk_region: str, is_s3: str, is_bi: str, kafka_config_bucket: str):
    aws_account_id = get_env_account(env_name)
    msk_env = get_msk_environment_name(env_name)
    bi_bucket = f'arn:aws:s3:::{kafka_config_bucket}'
    bi_bucket_all = f'arn:aws:s3:::{kafka_config_bucket}/*'
    return {
        **({"access_to_s3": {
            "Statement": [
                {
                    "Action": [
                        "s3:*"
                    ],
                    "Effect": "Allow",
                    "Resource": [
                        "*"
                    ]
                }
            ],
            "Version": "2012-10-17"
        }} if is_s3 else {}),
        f"{msk_env}-kafka-connector": {
            "Statement": [
                {
                    "Action": [
                        "kafka-cluster:*"
                    ],
                    "Effect": "Allow",
                    "Resource": get_kafka_resource_arn_list(env_name, aws_msk_region, aws_account_id)
                }
            ],
            "Version": "2012-10-17"
        },
        **({"s3-bi-connect": {
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Effect": "Allow",
                    "Action": [
                        "s3:ListAllMyBuckets"
                    ],
                    "Resource": "arn:aws:s3:::*"
                },
                {
                    "Effect": "Allow",
                    "Action": [
                        "s3:ListBucket",
                        "s3:GetBucketLocation"
                    ],
                    "Resource": [
                        bi_bucket
                    ]
                },
                {
                    "Effect": "Allow",
                    "Action": [
                        "s3:PutObject",
                        "s3:GetObject",
                        "s3:AbortMultipartUpload",
                        "s3:ListMultipartUploadParts",
                        "s3:ListBucketMultipartUploads",
                        "s3:PutObjectTagging"
                    ],
                    "Resource": [
                        bi_bucket_all
                    ]
                }
            ]
        }} if is_bi else {}),
        "cloudwatch-logs": {
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Effect": "Allow",
                    "Action": [
                        "logs:CreateLogGroup",
                        "logs:CreateLogStream",
                        "logs:PutLogEvents",
                        "logs:DescribeLogStreams",
                        "logs:DescribeLogGroups",
                        "logs:CreateLogDelivery",
                        "logs:GetLogDelivery",
                        "logs:DeleteLogDelivery",
                        "logs:ListLogDeliveries",
                        "logs:PutResourcePolicy",
                        "logs:DescribeResourcePolicies",
                    ],
                    "Resource": "*"
                }
            ]        
        }
    }

def main():
    is_s3 = os.getenv("KAFKA_CONNECT_ADD_S3_POLICY")
    aws_msk_region = get_required_env_var("AWS_MSK_REGION")
    is_bi = os.getenv("KAFKA_CONNECT_ADD_BI_POLICY")
    kafka_config_bucket = os.getenv("KAFKA_CONFIG_VAR_S3_BUCKET")
    kafka_connect_policy = os.getenv("KAFKA_CONNECT_POLICY")
    env_name = get_required_env_var('ENV_NAME')
    connector_name = get_required_env_var('CONNECTOR_NAME')

    iam_client = create_client_or_exit("iam", aws_msk_region)
    # Define the role name and policy
    role_name = get_msk_connector_role_name(env_name, connector_name)
    role_arn = get_or_create_role(role_name, iam_client)
    logger.debug(f"Role ARN: {role_arn}")
    write_github_env(role_arn, "KAFKA_CONNECTOR_ROLE_ARN")
    policies = get_msk_role_policies(env_name, aws_msk_region, is_s3, is_bi, kafka_config_bucket)
    logger.debug(f"Generated policies: {json.dumps(policies, indent=2)}")
    for policy, policy_document in policies.items():
        ensure_inline_policy(role_name, policy, policy_document, iam_client)
    # Attach Kafka Connect managed policy
    if kafka_connect_policy:
        ensure_managed_policy(role_name, kafka_connect_policy, iam_client)
    logger.info("IAM role and policies ensured successfully!")


if __name__ == "__main__":
    main()