# -------------------------------------------------------------------------
#  This file is part of the MindStudio project.
# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
#
# MindStudio is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#
#          http://license.coscl.org.cn/MulanPSL2
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
# -------------------------------------------------------------------------

from typing import Optional
from pydantic import BaseModel

from atk.configs.nodes_config import LOCAL_IP


class WorkerConfig(BaseModel):
    queues: Optional[str] = "atk_queue"
    pool: Optional[str] = "prefork"
    concurrency: Optional[int] = 4
    autoscale: Optional[str] = "4,4"
    cpu_numbers: Optional[int] = None
    prefetch_multiplier: Optional[int] = 16

    def __init__(self, ip=LOCAL_IP, port=9090, **data):
        super().__init__(**data)
        self.queues = f"{ip}:{port}@{self.get_name()}"

    def get_name(self):
        return "atk_queue"


class DatasetWorkerConfig(WorkerConfig):
    queues: Optional[str] = "atk_queue"

    def get_name(self):
        return "create_dataset"


class CpuRunWorkerConfig(WorkerConfig):

    def get_name(self):
        return "cpu_run"


class CompareWorkerConfig(WorkerConfig):

    def get_name(self):
        return "compare"


class DeviceRunWorkerConfig(WorkerConfig):
    pool: Optional[str] = "solo"
    concurrency: Optional[int] = 1
    autoscale: Optional[str] = "1,1"
    cpu_numbers: Optional[int] = 2

    def __init__(self, ip=LOCAL_IP, port=9090, device_id=0, **data):
        super().__init__(ip, port, **data)
        self.queues = f"{ip}:{port}@{self.get_name()}_{device_id}"

    def get_name(self):
        return "device_run"


class CleanOutputDataWorkerConfig(WorkerConfig):
    concurrency: Optional[int] = 1
    autoscale: Optional[str] = "1,1"

    def get_name(self):
        return "clean_output_data"


LOCAL_WORKERS = [
    DatasetWorkerConfig,
    DeviceRunWorkerConfig,
    CpuRunWorkerConfig,
    CompareWorkerConfig
]
REMOTE_WORKERS = [
    DatasetWorkerConfig,
    DeviceRunWorkerConfig
]