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

# Configuration
import yaml

with open("gateway_services.yaml", "r") as f:
    gateway_services = yaml.safe_load(f)

services = gateway_services["gateway_services"]
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_REQUEST_TIMEOUT = int(os.getenv("SERVICE_REQUEST_TIMEOUT", 300))
NODE_REQUEST_TIMEOUT = int(os.getenv("NODE_REQUEST_TIMEOUT", 30))


async def update_service_annotation(v1, service, external_ips):
    try:
        # Get the current service object
        service_name = service["name"]
        namespace = service["namespace"]
        service_obj = await v1.read_namespaced_service(service_name, namespace)

        # Check if the annotation needs to be updated
        current_annotation = service_obj.metadata.annotations.get(ANNOTATION_KEY)
        target_annotation = ",".join(external_ips) + "," + 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:
            # Check for external IP
            external_ips, node_names = [], []
            async for event in w.stream(
                v1.list_node,
                label_selector=NODE_LABEL,
                _request_timeout=NODE_REQUEST_TIMEOUT,
            ):
                node = event["object"]
                node_name = node.metadata.name

                for address in node.status.addresses:
                    if address.type == "ExternalIP":
                        external_ips.append(address.address)
                        node_names.append(node_name)
            if len(external_ips) > 0:
                print(
                    f"Detected external IPs {str(external_ips)} for node {str(node_names)}",
                    flush=True,
                )
                for service in services:
                    await update_service_annotation(v1, service, external_ips)

        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_service(service):
    config.load_incluster_config()
    v1 = client.CoreV1Api()
    w = watch.Watch()
    service_name = service["name"]
    namespace = service["namespace"]
    label_selector = service["label_pattern"]

    while True:
        try:
            async for event in w.stream(
                v1.list_namespaced_service,
                namespace,
                label_selector=label_selector,
                _request_timeout=SERVICE_REQUEST_TIMEOUT,
            ):
                service_obj = event["object"]
                if (
                    event["type"] == "ADDED"
                    and service_obj.metadata.name == service_name
                ):
                    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)
                    external_ips = []
                    for node in nodes.items:
                        for address in node.status.addresses:
                            if address.type == "ExternalIP":
                                external_ips.append(address.address)
                    await update_service_annotation(v1, service, external_ips)

        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
    service_watchers = [watch_service(service) for service in services]
    await asyncio.gather(watch_nodes(), *service_watchers)


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