import asyncio
from kubernetes_asyncio import client, config, watch
import os
import logging

# Configuration
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")
SERVICE_LABEL_KEY = os.getenv("SERVICE_LABEL_KEY", "zlanservice")
SERVICE_LABEL_VALUE = os.getenv("SERVICE_LABEL_VALUE", "true")
SERVICE_REQUEST_TIMEOUT = int(os.getenv("SERVICE_REQUEST_TIMEOUT", 300))
NODE_REQUEST_TIMEOUT = int(os.getenv("NODE_REQUEST_TIMEOUT", 30))

# Logging configuration
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


async def update_service_annotation(v1, service, external_ips):
    try:
        service_name = service.metadata.name
        namespace = service.metadata.namespace
        service_obj = await v1.read_namespaced_service(service_name, namespace)

        current_annotation = service_obj.metadata.annotations.get(ANNOTATION_KEY)
        target_annotation = ",".join(external_ips) + "," + ZERO_GATEWAY_IP
        if current_annotation != target_annotation:
            body = {"metadata": {"annotations": {ANNOTATION_KEY: target_annotation}}}
            await v1.patch_namespaced_service(service_name, namespace, body)
            logger.info(
                f"Updated service {service_name} with new external IP: {target_annotation}"
            )

    except client.exceptions.ApiException as e:
        logger.error(f"API Exception in update_service_annotation: {e}")


async def watch_nodes(v1, external_ips_update_queue):
    w = watch.Watch()
    while True:
        try:
            async for event in w.stream(
                v1.list_node,
                label_selector=NODE_LABEL,
                _request_timeout=NODE_REQUEST_TIMEOUT,
            ):
                node = event["object"]
                external_ips = [
                    addr.address
                    for addr in node.status.addresses
                    if addr.type == "ExternalIP"
                ]
                if external_ips:
                    await external_ips_update_queue.put(external_ips)

        except client.exceptions.ApiException as e:
            logger.error(f"API Exception in watch_nodes: {e}")
            await asyncio.sleep(5)
        except asyncio.CancelledError:
            logger.info("Watch task was cancelled.")
            break
        except Exception as e:
            logger.error(f"Unexpected error in watch_nodes: {e}")
            await asyncio.sleep(5)


async def watch_services(v1, external_ips_update_queue):
    w = watch.Watch()
    label_selector = f"{SERVICE_LABEL_KEY}={SERVICE_LABEL_VALUE}"

    while True:
        try:
            async for event in w.stream(
                v1.list_service_for_all_namespaces,
                label_selector=label_selector,
                _request_timeout=SERVICE_REQUEST_TIMEOUT,
            ):
                service = event["object"]
                if event["type"] in {"ADDED", "MODIFIED"}:
                    external_ips = await external_ips_update_queue.get()
                    await update_service_annotation(v1, service, external_ips)

        except client.exceptions.ApiException as e:
            logger.error(f"API Exception in watch_services: {e}")
            await asyncio.sleep(5)
        except asyncio.CancelledError:
            logger.info("Watch task was cancelled.")
            break
        except Exception as e:
            logger.error(f"Unexpected error in watch_services: {e}")
            await asyncio.sleep(5)


async def main():
    config.load_incluster_config()
    v1 = client.CoreV1Api()

    external_ips_update_queue = asyncio.Queue()
    await asyncio.gather(
        watch_nodes(v1, external_ips_update_queue),
        watch_services(v1, external_ips_update_queue),
    )


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