import time
from pathlib import Path
from typing import Any

from lite_dist2.common import float2hex, int2hex
from lite_dist2.config import TableConfig, TableConfigProvider, WorkerConfig
from lite_dist2.curriculum_models.study_portables import StudyRegistry
from lite_dist2.study_strategies import StudyStrategyModel
from lite_dist2.suggest_strategies import SuggestStrategyModel
from lite_dist2.suggest_strategies.base_suggest_strategy import SuggestStrategyParam
from lite_dist2.table_node_api.start_table_api import start_in_thread
from lite_dist2.table_node_api.table_param import StudyRegisterParam
from lite_dist2.type_definitions import RawParamType, RawResultType
from lite_dist2.value_models.aligned_space_registry import LineSegmentRegistry, ParameterAlignedSpaceRegistry
from lite_dist2.worker_node.table_node_client import TableNodeClient
from lite_dist2.worker_node.trial_runner import AutoMPTrialRunner
from lite_dist2.worker_node.worker import Worker


def set_table_config() -> None:
    config = TableConfig(
        port=8000,
        curriculum_path=Path(__file__).parent / "example_curriculum.json",
        curriculum_save_interval_seconds=10,
    )
    TableConfigProvider.set(config)


def register_study(table_ip: str, table_port: int) -> None:
    _resolution = 10
    _half_size = 2.0

    study_register_param = StudyRegisterParam(
        study=StudyRegistry(
            name="mandelbrot",
            required_capacity=set(),
            study_strategy=StudyStrategyModel(type="all_calculation", study_strategy_param=None),
            suggest_strategy=SuggestStrategyModel(
                type="sequential",
                suggest_strategy_param=SuggestStrategyParam(strict_aligned=True),
            ),
            result_type="scalar",
            result_value_type="int",
            parameter_space=ParameterAlignedSpaceRegistry(
                type="aligned",
                axes=[
                    LineSegmentRegistry(
                        name="x",
                        type="float",
                        size=int2hex(_resolution),
                        step=float2hex(2 * _half_size / _resolution),
                        start=float2hex(-1 * _half_size),
                    ),
                    LineSegmentRegistry(
                        name="y",
                        type="float",
                        size=int2hex(_resolution),
                        step=float2hex(2 * _half_size / _resolution),
                        start=float2hex(-1 * _half_size),
                    ),
                ],
            ),
        ),
    )
    client = TableNodeClient(table_ip, table_port)
    client.register_study(study_register_param)


class Mandelbrot(AutoMPTrialRunner):
    _ABS_THRESHOLD = 2.0
    _MAX_ITER = 255

    def func(self, parameters: RawParamType, *_: tuple[Any, ...], **__: dict[str, Any]) -> RawResultType:
        x = float(parameters[0])
        y = float(parameters[1])
        c = complex(x, y)
        z = complex(0, 0)
        iter_count = 0
        while abs(z) <= self._ABS_THRESHOLD and iter_count < self._MAX_ITER:
            z = z**2 + c
            iter_count += 1
        return iter_count


def run_worker(table_ip: str, table_port: int) -> None:
    worker_config = WorkerConfig(
        name="w_01",
        process_num=2,
        max_size=10,
        wait_seconds_on_no_trial=5,
        table_node_request_timeout_seconds=60,
    )
    worker = Worker(
        trial_runner=Mandelbrot(),
        ip=table_ip,
        port=table_port,
        config=worker_config,
    )
    worker.start(stop_at_no_trial=True)


if __name__ == "__main__":
    # 1. Set table config for example
    #    Execution node type in normal use: Table
    #    NOTE: For normal use, change the `table_config.json` file in the root directory.
    set_table_config()

    # 2. Start table node in other thread
    #    Execution node type in normal use: Table
    #    NOTE: In most cases, use `start()` (or `uv run start`) instead of `start_in_thread()`
    #          since table nodes and worker nodes are separated.
    table_thread = start_in_thread()
    time.sleep(1)  # wait for activate

    # 3. Register study to table node
    #    Execution node type in normal use: Management
    #    NOTE: If the management node is also a worker node or a table node,
    #          `TableNodeClient.register_study` is available.
    #          Otherwise, studies can be registered using the curl command.
    register_study(table_ip="127.0.0.1", table_port=8000)

    # 4. run worker node
    #    Execution node type in normal use: Worker
    #    NOTE: In this example, the entire process should be completed within 10 seconds.
    #          10 seconds later, example_curriculum.json should be saved in the same directory as this file.
    #          The saving interval can be changed from the `TableConfig.curriculum_save_interval_seconds`.
    run_worker(table_ip="127.0.0.1", table_port=8000)

    # 5. Stop table node
    #    Execution node type in normal use: Table
    #    NOTE: Used to explicitly terminate a table node.
    #          In this example, however, the node is running on a daemon thread,
    #          so there is no need to terminate it explicitly.
    # table_thread.stop()
    # table_thread.join()
