import json
from typing import Any, Dict, List

import boto3


def _iam_client(aws_profile: str):
    """Return an IAM client scoped to the provided profile (account id)."""

    session = boto3.Session(profile_name=aws_profile)
    return session.client("iam")


def _sns_client(aws_profile: str, region: str):
    """Return an SNS client scoped to the provided profile and region."""

    session = boto3.Session(profile_name=aws_profile, region_name=region)
    return session.client("sns")


def _normalize_federated(value: Any) -> List[str]:
    if value is None:
        return []
    if isinstance(value, list):
        return [str(v) for v in value]
    return [str(value)]


def _statement_has_oidc(statement: Dict[str, Any], oidc_issuer_arn: str) -> bool:
    principal = statement.get("Principal", {}) or {}
    federated = principal.get("Federated")
    federated_list = _normalize_federated(federated)
    return oidc_issuer_arn in federated_list


def insert_oidc_into_role(aws_profile: str, role_name: str, oidc_issuer: str, oidc_issuer_arn: str) -> bool:
    """Add the OIDC provider to the role trust policy if not already present."""

    iam = _iam_client(aws_profile)
    role = iam.get_role(RoleName=role_name)
    policy_doc = role["Role"]["AssumeRolePolicyDocument"]

    statements = policy_doc.get("Statement", []) or []
    if isinstance(statements, dict):
        statements = [statements]

    if any(_statement_has_oidc(stmt, oidc_issuer_arn) for stmt in statements):
        return False

    add_statement = {
        "Effect": "Allow",
        "Principal": {"Federated": oidc_issuer_arn},
        "Action": "sts:AssumeRoleWithWebIdentity",
        "Condition": {
            "StringLike": {
                f"{oidc_issuer}:sub": "system:serviceaccount:*:*",
                f"{oidc_issuer}:aud": "sts.amazonaws.com",
            }
        },
    }

    statements.append(add_statement)
    policy_doc["Statement"] = statements

    iam.update_assume_role_policy(RoleName=role_name, PolicyDocument=json.dumps(policy_doc))
    return True


def remove_oidc_from_role(aws_profile: str, role_name: str, oidc_issuer: str, oidc_issuer_arn: str) -> bool:
    """Remove the OIDC provider from the role trust policy if present."""

    iam = _iam_client(aws_profile)
    role = iam.get_role(RoleName=role_name)
    policy_doc = role["Role"]["AssumeRolePolicyDocument"]

    statements = policy_doc.get("Statement", []) or []
    if isinstance(statements, dict):
        statements = [statements]

    filtered = [stmt for stmt in statements if not _statement_has_oidc(stmt, oidc_issuer_arn)]
    if len(filtered) == len(statements):
        return False

    policy_doc["Statement"] = filtered
    iam.update_assume_role_policy(RoleName=role_name, PolicyDocument=json.dumps(policy_doc))
    return True


def insert_alarm_policy_into_sns(
    aws_profile: str,
    account_id: str,
    sns_topic_arn: str,
    cluster_name: str,
    region: str,
) -> bool:
    """Ensure the SNS topic policy allows CloudWatch alarms from the given account and region.

    Returns True if the policy was updated, False if no change was necessary.
    """

    sns = _sns_client(aws_profile, region)
    attrs = sns.get_topic_attributes(TopicArn=sns_topic_arn)
    policy_text = attrs.get("Attributes", {}).get("Policy")

    if not policy_text or policy_text == "None":
        policy_doc: Dict[str, Any] = {"Version": "2008-10-17", "Id": "", "Statement": []}
    else:
        policy_doc = json.loads(policy_text)

    statements = policy_doc.get("Statement", []) or []
    if isinstance(statements, dict):
        statements = [statements]

    sid = f"{account_id}-{cluster_name}-sns"
    new_statement = {
        "Sid": sid,
        "Effect": "Allow",
        "Principal": {"AWS": "*"},
        "Action": ["SNS:Publish"],
        "Resource": sns_topic_arn,
        "Condition": {
            "ArnLike": {"aws:SourceArn": f"arn:aws:cloudwatch:{region}:{account_id}:alarm:*"}
        },
    }

    existing = next((stmt for stmt in statements if stmt.get("Sid") == sid), None)
    if existing == new_statement:
        return False

    filtered = [stmt for stmt in statements if stmt.get("Sid") != sid]
    filtered.append(new_statement)

    policy_doc["Statement"] = filtered
    policy_doc["Id"] = sid

    sns.set_topic_attributes(
        TopicArn=sns_topic_arn,
        AttributeName="Policy",
        AttributeValue=json.dumps(policy_doc),
    )
    return True


def remove_alarm_policy_from_sns(
    aws_profile: str,
    account_id: str,
    sns_topic_arn: str,
    cluster_name: str,
    region: str,
) -> bool:
    """Remove the CloudWatch alarm publish permission from the SNS topic policy if present."""

    sns = _sns_client(aws_profile, region)
    attrs = sns.get_topic_attributes(TopicArn=sns_topic_arn)
    policy_text = attrs.get("Attributes", {}).get("Policy")

    if not policy_text or policy_text == "None":
        return False

    policy_doc = json.loads(policy_text)
    statements = policy_doc.get("Statement", []) or []
    if isinstance(statements, dict):
        statements = [statements]

    sid = f"{account_id}-{cluster_name}-sns"
    filtered = [stmt for stmt in statements if stmt.get("Sid") != sid]
    if len(filtered) == len(statements):
        return False

    if filtered:
        policy_doc["Statement"] = filtered
        sns.set_topic_attributes(
            TopicArn=sns_topic_arn,
            AttributeName="Policy",
            AttributeValue=json.dumps(policy_doc),
        )
    else:
        sns.set_topic_attributes(
            TopicArn=sns_topic_arn,
            AttributeName="Policy",
            AttributeValue="",
        )

    return True