from google.protobuf.json_format import MessageToDict
from ..proto import serialize_length_delimited, ComputeNodeFormat, parse_length_delimited
from ..node import Node
from .proto import (
    PermutiveWorkerConfiguration,
    SinkInput,
    ExportRole,
)

class DataSourcePermutive(Node):
    """
    Compute node that fetches a dataset from Permutive.
    """

    def __init__(
            self,
            name: str,
            credentials_dependency: str,
    ) -> None:
        config = PermutiveWorkerConfiguration(
            credentialsDependency=credentials_dependency,
            importRole={}
        )
        config_serialized = serialize_length_delimited(config)
        super().__init__(
            name,
            config=config_serialized,
            enclave_type="decentriq.permutive-worker",
            dependencies=[credentials_dependency],
            output_format=ComputeNodeFormat.RAW
        )

class PermutiveDataTransformer(Node):
    """
    Compute node that transforms data into a Permutive compliant form.
    """

    def __init__(
            self,
            name: str,
            credentials_dependency: str,
            input: SinkInput,
            import_id: str,
            segment_name: str,
            segment_code: str,
            input_has_headers: bool,
    ) -> None:
        config = PermutiveWorkerConfiguration(
            credentialsDependency=credentials_dependency,
            exportRole=ExportRole(input=input),
            importId = import_id,
            segmentName = segment_name,
            segmentCode = segment_code,
            inputHasHeaders=input_has_headers,
        )
        config_serialized = serialize_length_delimited(config)
        super().__init__(
            name,
            config=config_serialized,
            enclave_type="decentriq.permutive-worker",
            dependencies=[credentials_dependency, input.dependency],
            output_format=ComputeNodeFormat.RAW
        )

class PermutiveWorkerDecoder:
    def decode(self, config: bytes):
        config_decoded = PermutiveWorkerConfiguration()
        parse_length_delimited(config, config_decoded)
        return MessageToDict(config_decoded)

__all__ = [
    "DataSourcePermutive",
    "PermutiveDataTransformer",
    "PermutiveWorkerDecoder"
]
