import asyncio
import signal
import time
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path

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.value_models.const_param import ConstParam, ConstParamElement
from lite_dist2.worker_node.table_node_client import TableNodeClient
from lite_dist2.worker_node.trial_runner import SemiAutoMPTrialRunner
from lite_dist2.worker_node.worker import Worker


def worker_init() -> None:
    # Ignore KeyboardInterrupt on worker processes
    signal.signal(signal.SIGINT, signal.SIG_IGN)


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)


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

    study_register_param = StudyRegisterParam(
        study=StudyRegistry(
            name=f"mandelbrot-{int(time.time())}",
            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",
            const_param=ConstParam(
                consts=[
                    ConstParamElement(type="float", key="abs_threshold", value=float2hex(2.0)),
                    ConstParamElement(type="int", key="max_iter", value=int2hex(255)),
                ],
            ),
            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)
    await client.register_study(study_register_param)


class Mandelbrot(SemiAutoMPTrialRunner):
    def func(self, parameters: RawParamType, *_: object, **kwargs: object) -> RawResultType:
        abs_threshold = self.get_typed("abs_threshold", float, kwargs)
        max_iter = self.get_typed("max_iter", int, kwargs)
        x = float(parameters[0])
        y = float(parameters[1])
        c = complex(x, y)
        z = complex(0, 0)
        iter_count = 0
        while abs(z) <= abs_threshold and iter_count < max_iter:
            z = z**2 + c
            iter_count += 1
        return iter_count


async 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,
    )
    with ProcessPoolExecutor(max_workers=2, initializer=worker_init) as pool:
        worker = Worker(
            trial_runner=Mandelbrot(),
            ip=table_ip,
            port=table_port,
            pool=pool,
            config=worker_config,
        )
        try:
            await worker.start_async(stop_at_no_trial=True)
        except KeyboardInterrupt:
            pool.shutdown(wait=False, cancel_futures=True)


async def main() -> None:
    # 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.
    start_in_thread()
    await asyncio.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.
    await 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`.
    await run_worker(table_ip="127.0.0.1", table_port=8000)


if __name__ == "__main__":
    asyncio.run(main())
