import datetime
import importlib
import random
import subprocess

import pytest
import ray
import ray.exceptions

from geneva.runners.ray._mgr import ray_cluster
from geneva.runners.ray.raycluster import (
    CPU_ONLY_NODE,
    _HeadGroupSpec,
    _WorkerGroupSpec,
)


def test_cluster_startup() -> None:
    cluster_name = "geneva-integ-test"
    cluster_name += f"-{datetime.datetime.now().strftime('%Y-%m-%d-%H-%M')}"
    cluster_name += f"-{random.randint(0, 10000)}"

    with ray_cluster(
        name=cluster_name,
        namespace="geneva",
        # allocate at least a single worker so the test runs faster
        # that we save time on waiting for the actor to start
        worker_groups=[
            _WorkerGroupSpec(
                name="worker",
                min_replicas=1,
            )
        ],
        use_portforwarding=True,
    ):
        ray.get(ray.remote(lambda: 1).remote())


def test_cluster_startup_no_account() -> None:
    """
    Test the if we try to import geneva, which uses gcs for
    workspace packaging magic, errors when the service account
    doesn't have permission to access the gcs bucket.
    """

    cluster_name = "geneva-integ-test"
    cluster_name += f"-{datetime.datetime.now().strftime('%Y-%m-%d-%H-%M')}"
    cluster_name += f"-{random.randint(0, 10000)}"

    with (
        pytest.raises(ValueError, match=r"Service account .* does not exist"),
        ray_cluster(
            name=cluster_name,
            namespace="geneva",
            use_portforwarding=True,
            # allocate at least a single worker so the test runs faster
            # that we save time on waiting for the actor to start
            worker_groups=[
                _WorkerGroupSpec(
                    name="worker",
                    min_replicas=1,
                    service_account="bogus-service-account",
                )
            ],
        ),
    ):
        ray.get(ray.remote(lambda: importlib.import_module("geneva")).remote())


@pytest.mark.skip("Need to create a service account with no permissions")
def test_cluster_startup_no_gcs_permission() -> None:
    """
    Test the if we try to import geneva, which uses gcs for
    workspace packaging magic, errors when the service account
    doesn't have permission to access the gcs bucket.
    """

    cluster_name = "geneva-integ-test"
    cluster_name += f"-{datetime.datetime.now().strftime('%Y-%m-%d-%H-%M')}"
    cluster_name += f"-{random.randint(0, 10000)}"

    with (
        ray_cluster(
            name=cluster_name,
            namespace="geneva",
            use_portforwarding=True,
            # allocate at least a single worker so the test runs faster
            # that we save time on waiting for the actor to start
            worker_groups=[
                _WorkerGroupSpec(
                    name="worker",
                    min_replicas=1,
                    service_account="valid-no-perms-service-account",
                )
            ],
        ),
        pytest.raises(PermissionError, match="PERMISSION_DENIED"),
    ):
        ray.get(ray.remote(lambda: importlib.import_module("geneva")).remote())


def test_cluster_startup_can_import_geneva_and_lance(
    geneva_k8s_service_account: str,
) -> None:
    """
    Test the if we try to import geneva, which uses gcs for
    workspace packaging magic, errors when the service account
    doesn't have permission to access the gcs bucket.
    """

    cluster_name = "geneva-integ-test"
    cluster_name += f"-{datetime.datetime.now().strftime('%Y-%m-%d-%H-%M')}"
    cluster_name += f"-{random.randint(0, 10000)}"

    with ray_cluster(
        name=cluster_name,
        namespace="geneva",
        use_portforwarding=True,
        head_group=_HeadGroupSpec(
            service_account=geneva_k8s_service_account,
        ),
        # allocate at least a single worker so the test runs faster
        # that we save time on waiting for the actor to start
        worker_groups=[
            _WorkerGroupSpec(
                name="worker",
                min_replicas=1,
                service_account=geneva_k8s_service_account,
            )
        ],
    ):
        ray.get(ray.remote(lambda: importlib.import_module("geneva")).remote())
        ray.get(ray.remote(lambda: importlib.import_module("lance")).remote())


def test_cluster_startup_skip_site(
    geneva_k8s_service_account: str,
) -> None:
    """
    Test that if we try to import geneva, which uses gcs for
    workspace packaging magic, errors because we skipped site
    packages and can't find
    """

    cluster_name = "geneva-integ-test"
    cluster_name += f"-{datetime.datetime.now().strftime('%Y-%m-%d-%H-%M')}"
    cluster_name += f"-{random.randint(0, 10000)}"

    with (
        ray_cluster(
            name=cluster_name,
            namespace="geneva",
            use_portforwarding=True,
            head_group=_HeadGroupSpec(
                service_account=geneva_k8s_service_account,
            ),
            # allocate at least a single worker so the test runs faster
            # that we save time on waiting for the actor to start
            worker_groups=[
                _WorkerGroupSpec(
                    name="worker",
                    min_replicas=1,
                    service_account=geneva_k8s_service_account,
                )
            ],
            skip_site_packages=True,
        ),
        pytest.raises(ImportError, match=r"^(?!.*(geneva)).+"),
    ):
        # match the word geneva is not in the error message
        # so we know that geneva is uploaded but not lance, which
        # is expected because we skipped site packages
        ray.get(ray.remote(lambda: importlib.import_module("geneva")).remote())


def test_cluster_startup_cpu_only_tag(
    geneva_k8s_service_account: str,
) -> None:
    """
    Test that if we try to import geneva, which uses gcs for
    workspace packaging magic, errors when the service account
    doesn't have permission to access the gcs bucket.
    """

    cluster_name = "geneva-integ-test"
    cluster_name += f"-{datetime.datetime.now().strftime('%Y-%m-%d-%H-%M')}"
    cluster_name += f"-{random.randint(0, 10000)}"

    with ray_cluster(
        name=cluster_name,
        namespace="geneva",
        use_portforwarding=True,
        head_group=_HeadGroupSpec(
            service_account=geneva_k8s_service_account,
        ),
        # allocate both GPU and CPU workers
        worker_groups=[
            _WorkerGroupSpec(
                name="cpu",
                min_replicas=0,
                service_account=geneva_k8s_service_account,
                num_cpus=2,
            ),
            _WorkerGroupSpec(
                name="gpu",
                min_replicas=1,
                service_account=geneva_k8s_service_account,
                num_gpus=1,
                num_cpus=2,
            ),
        ],
    ):
        ray.get(
            ray.remote(num_gpus=1)(
                lambda: subprocess.run("nvidia-smi", check=True)
            ).remote()
        )
        # running a CPU only task but we have a GPU worker
        # should schedule on the GPU worker
        ray.get(ray.remote(lambda: subprocess.run("nvidia-smi", check=True)).remote())

        # running a CPU only task + force it to run on CPU
        # should schedule on the CPU worker
        with pytest.raises(
            ray.exceptions.RayTaskError,
            match="No such file or directory: 'nvidia-smi'",
        ):
            ray.get(
                ray.remote(resources={CPU_ONLY_NODE: 1})(
                    lambda: subprocess.run("nvidia-smi", check=True)
                ).remote()
            )
