from typing import List, Optional, Dict, Set, Any
from uuid import UUID

from pydantic import validator

from unipipeline.modules.uni_definition import UniDefinition
from unipipeline.modules.uni_broker_definition import UniBrokerDefinition
from unipipeline.modules.uni_message_definition import UniMessageDefinition
from unipipeline.modules.uni_module_definition import UniModuleDefinition
from unipipeline.modules.uni_waiting_definition import UniWaitingDefinition


class UniWorkerDefinition(UniDefinition):
    id: UUID
    name: str
    broker: UniBrokerDefinition[Any]
    type: UniModuleDefinition
    topic: str
    error_topic: str
    error_payload_topic: str
    input_message: UniMessageDefinition
    output_workers: Set[str]
    ack_after_success: bool
    waitings: Set[UniWaitingDefinition]

    def get_related_broker_names(self, index_of_workers: Dict[str, 'UniWorkerDefinition']) -> Set[str]:
        result = {self.broker.name,}
        for w_name in self.output_workers:
            result.add(index_of_workers[w_name].broker.name)
        return result

    def get_related_worker_names(self) -> Set[str]:
        result = {self.name,}
        for w_name in self.output_workers:
            result.add(w_name)
        return result
