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

# 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 = os.getenv("SERVICE_LABEL", "enablezlan=true")
ZLAN_GATEWAY_IP_KEY = os.getenv("ZLAN_GATEWAY_IP_KEY", "zlanip")
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.DEBUG)
logger = logging.getLogger(__name__)


def is_valid_ip(ip):
    try:
        ipaddress.ip_address(ip)
        return True
    except ValueError:
        return False


async def update_service_annotation(v1, service, external_ips):
    try:
        service_name = service.metadata.name
        namespace = service.metadata.namespace
        logger.debug(f"Fetching service {service_name} in namespace {namespace}")
        service_obj = await v1.read_namespaced_service(service_name, namespace)

        current_annotation = service_obj.metadata.annotations.get(ANNOTATION_KEY)
        zlan_gateway_ip = service_obj.metadata.labels.get(ZLAN_GATEWAY_IP_KEY)
        logger.debug(f"Zlan Gateway IP: {zlan_gateway_ip}")

        if is_valid_ip(zlan_gateway_ip):
            target_annotation = ",".join(external_ips) + "," + zlan_gateway_ip
        else:
            target_annotation = ",".join(external_ips)
            logger.debug(
                f"Invalid Zlan Gateway IP: {zlan_gateway_ip}, excluding from target annotation"
            )

        logger.debug(f"Current annotation: {current_annotation}")
        logger.debug(f"Target annotation: {target_annotation}")
        if current_annotation != target_annotation:
            body = {"metadata": {"annotations": {ANNOTATION_KEY: target_annotation}}}
            logger.debug(f"Patching service {service_name} with body: {body}")
            await v1.patch_namespaced_service(service_name, namespace, body)
            logger.info(
                f"Updated service {service_name} with new external IP: {target_annotation}"
            )
        else:
            logger.debug(f"No update required for service {service_name}")

    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()
    external_node_ipset = set()
    while True:
        try:
            logger.debug("Starting to watch nodes")
            async for event in w.stream(
                v1.list_node,
                label_selector=NODE_LABEL,
                _request_timeout=NODE_REQUEST_TIMEOUT,
            ):
                node = event["object"]
                event_type = event["type"]
                logger.debug(
                    f"Received {event_type} event for node: {node.metadata.name}"
                )
                external_ips = [
                    addr.address
                    for addr in node.status.addresses
                    if addr.type == "ExternalIP"
                ]
                if event_type in {"ADDED", "MODIFIED"}:
                    external_node_ipset.update(external_ips)
                    logger.debug(
                        f"External IPs for node {node.metadata.name}: {external_ips}"
                    )
                elif event_type == "REMOVED":
                    for ip in external_ips:
                        if ip in external_node_ipset:
                            external_node_ipset.remove(ip)
                    logger.debug(
                        f"External IPs for node {node.metadata.name}: {external_ips}"
                    )
                external_node_ips = list(external_node_ipset)
                await external_ips_update_queue.put(external_node_ips)
                logger.debug(f"Added external IPs to update queue: {external_node_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()

    while True:
        try:
            logger.debug("Starting to watch services")
            async for event in w.stream(
                v1.list_service_for_all_namespaces,
                label_selector=SERVICE_LABEL,
                _request_timeout=SERVICE_REQUEST_TIMEOUT,
            ):
                service = event["object"]
                if event["type"] in {"ADDED", "MODIFIED"}:
                    logger.debug(
                        f"Processing event type: {event['type']} for service: {service.metadata.name}"
                    )
                    while True:
                        # maintain only last set of external_ips by popping everything
                        # and putting back the final one when it runs empty
                        external_ips = await external_ips_update_queue.get()
                        if external_ips_update_queue.empty():
                            await external_ips_update_queue.put(external_ips)
                            break
                    logger.debug(f"Retrieved external IPs: {external_ips}")
                    await update_service_annotation(v1, service, external_ips)
                    logger.debug(f"Service updated: {service.metadata.name}")

        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())