import pytest

from ldnode.calculator.handlers import handler_registry
from ldnode.calculator.runtime_config import (
    calculate_hyper_search_runtime_resources,
    configure_hyper_search_runtime_resources,
    get_hyper_search_runtime_resources,
)


@pytest.fixture(autouse=True)
def restore_hyper_search_runtime_resources():
    hyper_search_handler = handler_registry.get("hyper_search")
    original_cpu_usage = hyper_search_handler.cpu_usage
    original_resources = get_hyper_search_runtime_resources()

    yield

    hyper_search_handler.cpu_usage = original_cpu_usage
    configure_hyper_search_runtime_resources(original_resources.cpu_resource)


def test_hyper_search_runtime_resources_leave_one_cpu_for_optuna_scheduler():
    resources = calculate_hyper_search_runtime_resources(7)

    assert resources.cpu_resource == 7
    assert resources.cpu_usage == 7
    assert resources.optuna_n_jobs == 6


def test_hyper_search_runtime_resources_keep_one_optuna_worker_for_small_node():
    resources = calculate_hyper_search_runtime_resources(1)

    assert resources.cpu_resource == 1
    assert resources.cpu_usage == 1
    assert resources.optuna_n_jobs == 1


def test_configure_hyper_search_runtime_resources_updates_runtime_config():
    configure_hyper_search_runtime_resources(7)

    resources = get_hyper_search_runtime_resources()

    assert resources.cpu_resource == 7
    assert resources.cpu_usage == 7
    assert resources.optuna_n_jobs == 6


def test_handler_registry_configures_hyper_search_cpu_usage_from_cpu_resource():
    handler_registry.configure_runtime_resources(7)

    assert handler_registry.get("hyper_search").cpu_usage == 7
