import asyncio

from lightkube import AsyncClient
from lightkube.resources.core_v1 import Node
from lightkube.types import PatchType


async def set_feature_label(
    client: AsyncClient, node_name: str, kind: str, variant: str
):
    """Add label with feature to the node"""
    # create patch
    patch_metadata = [
        {
            "op": "add",
            "path": f"/metadata/labels/node-feature.k8s.dpc24.ru~1{kind}",
            "value": variant,
        }
    ]
    # run patch
    await client.patch(
        Node, name=node_name, obj=patch_metadata, patch_type=PatchType.JSON
    )


async def unset_feature_label(
    client: AsyncClient, node_name: str, kind: str, variant: str
):
    """Remove label with feature to the node"""
    # delete patch
    patch_metadata = [
        {
            "op": "remove",
            "path": f"/metadata/labels/node-feature.k8s.dpc24.ru~1{kind}",
        }
    ]
    # run patch
    await client.patch(
        Node, name=node_name, obj=patch_metadata, patch_type=PatchType.JSON
    )

async def main():
    lk = AsyncClient()
    await set_feature_label(lk, "feature-core-k8s2-w9-dpc", "test", "nvidia")


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