import codecs
import ipaddress
from typing import TYPE_CHECKING

import boto3
import sentry_sdk
from aws_lambda_powertools.utilities.parser import event_parser
from aws_lambda_powertools.utilities.typing import LambdaContext
from lambdacommon.common_config import logger
from sentry_sdk.integrations.aws_lambda import AwsLambdaIntegration

import update_waf_ipsets.config as Config
from update_waf_ipsets.models.payload import LambdaPayload, ScopeLiteral

if TYPE_CHECKING:
    from mypy_boto3_wafv2 import WAFV2Client

if Config.SENTRY_DSN:
    sentry_sdk.init(
        dsn=Config.SENTRY_DSN, environment=Config.ENVIRONMENT, integrations=[AwsLambdaIntegration(timeout_warning=True)]
    )

sts_client = boto3.client("sts")
s3_client = boto3.client("s3")


def assume_source_account_role(account_id: str, role_name: str, external_id: str) -> dict:
    """
    Assume IAM role in source account.

    Args:
    account_id (str): AWS account ID
    role_name (str): IAM role name to assume
    kwargs (dict): Additional parameters for assume_role

    Returns:
    dict: Dict of AWS credentials
    """

    response = sts_client.assume_role(
        RoleArn=f"arn:aws:iam::{account_id}:role/{role_name}",
        RoleSessionName=Config.ROLE_SESSION_NAME,
        DurationSeconds=Config.ROLE_SESSION_DURATION,
        ExternalId=external_id,
    )
    credentials = {}
    credentials["aws_access_key_id"] = response["Credentials"]["AccessKeyId"]
    credentials["aws_secret_access_key"] = response["Credentials"]["SecretAccessKey"]
    credentials["aws_session_token"] = response["Credentials"]["SessionToken"]

    return credentials


def get_ip_addresses(bucket: str, key: str) -> list[str]:
    """Gets and validates list of IP addresses from an S3 file.

    Args:
        bucket (str):
        key (str):

    Returns:
        list[str]:

    """
    s3_object_body = s3_client.get_object(Bucket=bucket, Key=key)["Body"]
    line_stream = codecs.getreader("utf-8")
    addresses = [line.strip() for line in line_stream(s3_object_body)]
    for address in addresses:
        assert ipaddress.ip_network(address), "Not a valid network!"
    return addresses


def get_ip_set_by_name(client: "WAFV2Client", name: str, scope: ScopeLiteral) -> dict:
    """Get an IP Set by name.

    Args:
        client (boto3.client): boto3 client for WAFv2
        name (str): Name of the IP Set
        scope (ScopeLiteral): Scope of the IP Set (CLOUDFRONT or REGIONAL)

    Returns:
        dict: IP Set details or empty dict if not found
    """
    response = client.list_ip_sets(Scope=scope)
    for ip_set in response["IPSets"]:
        if ip_set["Name"] == name:
            return ip_set
    logger.error(f"IP Set {name} is not found.")
    raise ValueError(f"IP Set {name} is not found.")


def update_waf_ip_set(client: "WAFV2Client", addresses: list[str], scope: ScopeLiteral, name: str) -> None:
    """Updates waf ip set

    Args:
        client (boto3.client): boto3 client for WAFv2
        addresses (list[str]): list of IP addresses to update ip set with
        scope (ScopeLiteral): CLOUDFRONT or REGIONAL scope
        name (str): name of waf ip set to update

    Returns:
        None
    """

    matching_ip_set = get_ip_set_by_name(client, name, scope)
    logger.info(f"Found {matching_ip_set}")
    kwargs = {}
    if "Description" in matching_ip_set:
        if matching_ip_set["Description"] != "":
            kwargs["Description"] = matching_ip_set["Description"]
    response = client.update_ip_set(
        Name=matching_ip_set["Name"],
        Scope=scope,
        Id=matching_ip_set["Id"],
        LockToken=matching_ip_set["LockToken"],
        Addresses=addresses,
        **kwargs,
    )
    logger.info(f"IP Set {name} updated")
    logger.debug(f"Response: {response}")


@event_parser(model=LambdaPayload)
def handler(event: LambdaPayload, context: LambdaContext) -> dict:
    # Read new ip lists into memory from s3
    logger.info(f"Getting ip list from s3://{event.s3_bucket}/{event.s3_key}...")
    addresses = get_ip_addresses(event.s3_bucket, event.s3_key)

    logger.info(
        f"Assuming role {event.role_name} in account {event.account_id} using external id {event.external_id}..."
    )
    credentials = assume_source_account_role(
        event.account_id,
        event.role_name,
        event.external_id,
    )

    for ip_set in event.ip_sets:
        # Override region to default for cloudfront-scoped ip sets
        if ip_set.scope == "CLOUDFRONT":
            ip_set.region = Config.DEFAULT_REGION
        logger.info(f"Updating IP set {ip_set.name}, scope: {ip_set.scope}, region: {ip_set.region} ")
        waf_client = boto3.client("wafv2", region_name=ip_set.region, **credentials)
        update_waf_ip_set(waf_client, addresses, ip_set.scope, ip_set.name)

    return {
        "accound_id": event.account_id,
        "ip_sets": [ip_set.model_dump() for ip_set in event.ip_sets],
        "result": "SUCCESS",
    }