from pyinfra import host
from pyinfra.api import deploy
from pyinfra.facts import files
from pyinfra.operations import files as files_op
from pyinfra.operations import server

K3S_COMMON_FLAGS = "--disable local-storage --disable=servicelb --disable=traefik"
IS_CREATED_FILE = "/pyinfra/.k3s-init-completed"

@deploy("Deploy K3S")
def deploy_k3s():
    k3s_config = host.data.get("k3s")

    if not k3s_config:
        raise ValueError("k3s ")

    server.shell(
        name="Disable ufw",
        commands=["ufw disable"],
    )

    if host.get_fact(files.File, IS_CREATED_FILE):
        return

    if k3s_config.get("role") == "master":
        server.shell(
            name="Start K3S master",
            commands=[
                "curl -sfL https://get.k3s.io | "
                f"sh -s - server --cluster-init {K3S_COMMON_FLAGS}",
            ],
            _env=dict(K3S_TOKEN=k3s_config.get("token")),
        )
    else:
        server.shell(
            name="Start K3S worker",
            commands=[
                "curl -sfL https://get.k3s.io | "
                f"sh -s - server --server https://{k3s_config.get('master')}:6443 {K3S_COMMON_FLAGS}",
            ],
            _env=dict(K3S_TOKEN=k3s_config.get("token")),
        )
    files_op.file(
        name="Creating persistence to not run this again",
        path=IS_CREATED_FILE,
        present=True,
    )