import threading as th
from typing import Any, Callable, Generator

from ldnode.calculator.operators import (
    HyperparameterOptimizer,
    HyperparameterSearchStorage,
)
from ldnode.calculator.schemas import CancellationToken

from .models import HyperparameterSearchParameters, TaskResponse
from .utils import Collector, wrap_handler


@wrap_handler("超参数搜索", is_generator=True)
def handle_hyperparameter_search(
    parameters: HyperparameterSearchParameters,
    invoke: Callable[[str, dict[str, Any], float], Any],
    cancel_token: CancellationToken,
) -> Generator[TaskResponse[Any]]:
    """响应超参数搜索任务"""
    collector = Collector()
    storage = HyperparameterSearchStorage(parameters.get("history_results"))
    optimizer = HyperparameterOptimizer(
        study_id=parameters["study_id"],
        method=parameters["method"],
        config=parameters["config"],
        parameters=parameters["parameters"],
        test_begin_date=parameters.get("test_begin_date"),
        storage=storage,
        invoke_handler=invoke,
        cancel_token=cancel_token,
        result_collector=collector,
    )
    runner = th.Thread(target=optimizer.search, daemon=True)
    runner.start()
    for result in collector.read():
        yield result
    runner.join()
