import boto3, socket, os

elbv2 = boto3.client("elbv2")

def resolve_ip(hostname):
    return list({res[4][0] for res in socket.getaddrinfo(hostname, None)})

def lambda_handler(event, context):
    rds_endpoint = os.environ["RDS_PROXY_ENDPOINT"]
    tg_arn       = os.environ["TARGET_GROUP_ARN"]
    port         = int(os.environ["TARGET_PORT"])

    ips = resolve_ip(rds_endpoint)
    print(f"Resolved IPs for {rds_endpoint}: {ips}")

    # 当前已注册的 targets
    current_targets = elbv2.describe_target_health(TargetGroupArn=tg_arn)
    current_ips = [t["Target"]["Id"] for t in current_targets["TargetHealthDescriptions"]]
    print(f"Currently registered: {current_ips}")

    # 注册新 IP
    for ip in ips:
        if ip not in current_ips:
            print(f"Registering {ip}")
            elbv2.register_targets(
                TargetGroupArn=tg_arn,
                Targets=[{"Id": ip, "Port": port}]
            )

    # 移除无效 IP
    for ip in current_ips:
        if ip not in ips:
            print(f"Deregistering {ip}")
            elbv2.deregister_targets(
                TargetGroupArn=tg_arn,
                Targets=[{"Id": ip, "Port": port}]
            )

    return {"status": "done", "new_ips": ips}