import codecs

import boto3
from lambdacommon.common_config import logger

from update_dns_firewall.models.payload import DomainListPayload


def get_domains(bucket, key):
    """
    Gets list of domains from an S3 file.
    """
    s3 = boto3.client("s3")
    obj = s3.get_object(Bucket=bucket, Key=key)
    lines = codecs.getreader("utf-8")(obj["Body"])
    domains = [line.strip() for line in lines if line.strip()]
    return domains


def chunks(iterable, size):
    """
    Splits the domain into a specified size.
    """
    for i in range(0, len(iterable), size):
        yield iterable[i : i + size]


def update_firewall_domain_list(domain_list_id, domains, region):
    """
    Updates DNS Firewall domain list with the provided domains.
    """
    client = boto3.client("route53resolver", region_name=region)
    domain_chunks = list(chunks(domains, 1000))
    for idx, chunk in enumerate(domain_chunks):
        op = "REPLACE" if idx == 0 else "ADD"
        client.update_firewall_domains(FirewallDomainListId=domain_list_id, Operation=op, Domains=chunk)
        print(f"{op}: {len(chunk)} domains ({idx * 1000 + 1}-{idx * 1000 + len(chunk)})")


def handler(event, context):
    try:
        payload = DomainListPayload(**event)
        logger.info(f"Getting domain list from s3://{payload.s3_bucket}/{payload.s3_key}...")
        domains = get_domains(payload.s3_bucket, payload.s3_key)
        logger.info(f"Retrieved {len(domains)} domains from S3.")
        update_firewall_domain_list(payload.domain_list_id, domains, payload.region)
        logger.info("DNS Firewall domain list update complete.")

        return {
            "domain_length": len(domains),
            "result": "SUCCESS",
        }
    except Exception as e:
        logger.exception(f"Error updating DNS Firewall domain list: {e}")
        raise e