from kubernetes import client, config, watch
import os

# Load in-cluster config
config.load_incluster_config()

# Set up Kubernetes API client
v1 = client.CoreV1Api()

# Configuration
SERVICE_NAME = "traefik"
NAMESPACE = "kube-system"
ANNOTATION_KEY = "kube-vip.io/loadbalancerIPs"


def update_service_annotation(external_ip):
    # Get the current service object
    service = v1.read_namespaced_service(SERVICE_NAME, NAMESPACE)

    # Check if the annotation needs to be updated
    current_annotation = service.metadata.annotations.get(ANNOTATION_KEY)
    if current_annotation != external_ip:
        # Update the annotation
        body = {"metadata": {"annotations": {ANNOTATION_KEY: external_ip}}}
        v1.patch_namespaced_service(SERVICE_NAME, NAMESPACE, body)
        print(f"Updated service {SERVICE_NAME} with new external IP: {external_ip}")


def main():
    w = watch.Watch()
    for event in w.stream(v1.list_node, _request_timeout=60):
        node = event["object"]
        node_name = node.metadata.name

        # Extract the external IP if it exists
        external_ip = None
        for address in node.status.addresses:
            if address.type == "ExternalIP":
                external_ip = address.address
                break

        if external_ip:
            print(f"Detected external IP {external_ip} for node {node_name}")
            update_service_annotation(external_ip)


if __name__ == "__main__":
    main()