import asyncio
from kubernetes_asyncio import client, config, watch

# node_external_ip_controller_async.py

import os

# Configuration
SERVICE_NAME = os.getenv("SERVICE_NAME", "traefik")
SERVICE_NAME_LABEL_PATTERN = os.getenv(
    "SERVICE_NAME_LABEL_PATTERN", "app.kubernetes.io/name=traefik"
)
NAMESPACE = os.getenv("NAMESPACE", "kube-system")
ANNOTATION_KEY = os.getenv("ANNOTATION_KEY", "kube-vip.io/loadbalancerIPs")
ZERO_GATEWAY_IP = os.getenv("ZERO_GATEWAY_IP", "172.28.10.1")
NODE_LABEL = os.getenv("NODE_LABEL", "svccontroller.k3s.cattle.io/enablelb=true")


async def update_service_annotation(v1, service_name, external_ip):
    try:
        # Get the current service object
        service = await v1.read_namespaced_service(service_name, NAMESPACE)

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

    except client.exceptions.ApiException as e:
        print(f"API Exception in update_service_annotation: {e}", flush=True)


async def watch_nodes():
    config.load_incluster_config()
    v1 = client.CoreV1Api()
    w = watch.Watch()

    while True:
        try:
            async for event in w.stream(
                v1.list_node, label_selector=NODE_LABEL, _request_timeout=300
            ):
                node = event["object"]
                node_name = node.metadata.name

                # Check for external IP
                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}",
                        flush=True,
                    )

                    # Get all services that need to be updated with this external IP
                    # services = await v1.list_namespaced_service(
                    #     NAMESPACE, label_selector=SERVICE_NAME_PATTERN
                    # )
                    service = await v1.read_namespaced_service(SERVICE_NAME, NAMESPACE)
                    # for service in services.items:
                    await update_service_annotation(
                        v1, service.metadata.name, external_ip
                    )

        except client.exceptions.ApiException as e:
            print(f"API Exception in watch_nodes: {e}", flush=True)
            await asyncio.sleep(5)

        except asyncio.CancelledError:
            print("Watch task was cancelled.", flush=True)
            break

        except Exception as e:
            print(f"Unexpected error in watch_nodes: {e}", flush=True)
            await asyncio.sleep(5)


async def watch_services():
    config.load_incluster_config()
    v1 = client.CoreV1Api()
    w = watch.Watch()

    while True:
        try:
            async for event in w.stream(
                v1.list_namespaced_service,
                NAMESPACE,
                label_selector=SERVICE_NAME_LABEL_PATTERN,
                _request_timeout=300,
            ):
                service = event["object"]
                service_name = service.metadata.name
                if event["type"] == "ADDED":
                    print(f"New service detected: {service_name}", flush=True)

                    # Fetch current external IP for nodes that match the label
                    nodes = await v1.list_node(label_selector=NODE_LABEL)
                    for node in nodes.items:
                        for address in node.status.addresses:
                            if address.type == "ExternalIP":
                                external_ip = address.address
                                await update_service_annotation(
                                    v1, service_name, external_ip
                                )
                                break

        except client.exceptions.ApiException as e:
            print(f"API Exception in watch_services: {e}", flush=True)
            await asyncio.sleep(5)

        except asyncio.CancelledError:
            print("Watch task was cancelled.", flush=True)
            break

        except Exception as e:
            print(f"Unexpected error in watch_services: {e}", flush=True)
            await asyncio.sleep(5)


async def main():

    # Run watch_nodes and watch_services concurrently
    await asyncio.gather(watch_nodes(), watch_services())


if __name__ == "__main__":
    asyncio.run(main())