__init__.py

import importlib
import os
import sys
import warnings
from enum import Enum
from pathlib import Path
from typing import Union

import google.protobuf.internal.api_implementation
from google.protobuf import symbol_database as _symbol_database

from snet.sdk.exceptions import NoGroupsFoundError, GroupNotFoundError, ServiceMetadataMismatchError
from snet.sdk.registry.models import StorageType, FileURI
from snet.sdk.registry.organization_metadata import OrganizationMetadata
from snet.sdk.registry.registry_contract import RegistryContract
from snet.sdk.registry.service_metadata import ServiceMetadata, Group

with warnings.catch_warnings():
    # Suppress the eth-typing package`s warnings related to some new networks
    warnings.filterwarnings(
        "ignore",
        "Network .* does not have a valid ChainId. eth-typing should be "
        "updated with the latest networks.",
        UserWarning,
    )

from snet.sdk.account import Account
from snet.sdk.config import config
from snet.sdk.client_lib_generator import ClientLibGenerator
from snet.sdk.mpe.mpe_contract import MPEContract
from snet.sdk.mpe.payment_channel_provider import PaymentChannelProvider
from snet.sdk.payment_strategies import (
    DefaultPaymentStrategy,
    PaidCallPaymentStrategy,
    PrePaidPaymentStrategy,
    FreeCallPaymentStrategy,
    PaymentStrategy,
)
from snet.sdk.service_client import ServiceClient
from snet.sdk.registry.storage_provider import StorageProvider
from snet.sdk.types import ModuleName, ServiceStub
from snet.sdk.utils.utils import (
    find_file_by_keyword,
    get_we3_object,
)

google.protobuf.internal.api_implementation.Type = lambda: "python"
_sym_db = _symbol_database.Default()
_sym_db.RegisterMessage = lambda x: None


class PaymentStrategyType(Enum):
    PAID_CALL = PaidCallPaymentStrategy
    FREE_CALL = FreeCallPaymentStrategy
    PREPAID_CALL = PrePaidPaymentStrategy
    DEFAULT = DefaultPaymentStrategy


class SnetSDK:
    def __init__(self):
        self.w3 = get_we3_object()
        self.mpe_contract = MPEContract()
        self.registry_contract = RegistryContract()
        self.storage_provider = StorageProvider()
        self.payment_channel_provider = PaymentChannelProvider(self.mpe_contract)
        self.account = Account(self.mpe_contract.contract.address)

    def create_service_client(
        self,
        org_id: str,
        service_id: str,
        group_name: str = None,
        payment_strategy: PaymentStrategy = None,
        payment_strategy_type: PaymentStrategyType = PaymentStrategyType.DEFAULT,
        options=None,
        concurrent_calls: int = 1,
    ):
        service_metadata = self._enhance_service_metadata(org_id, service_id)
        lib_generator = ClientLibGenerator(self.storage_provider, org_id, service_id)

        if service_metadata.service_api_source is not None:
            service_api_source = service_metadata.service_api_source
        else:
            service_api_source = service_metadata.model_ipfs_hash
        service_api_source = FileURI.from_raw_uri(service_api_source)

        force_update = config.FORCE_UPDATE
        if force_update:
            lib_generator.generate_client_library(service_api_source)
        else:
            path_to_pb_files = lib_generator.proto_dir
            pb_2_file_name = find_file_by_keyword(
                path_to_pb_files, keyword="pb2.py", exclude=["training"]
            )
            pb_2_grpc_file_name = find_file_by_keyword(
                path_to_pb_files, keyword="pb2_grpc.py", exclude=["training"]
            )
            if not pb_2_file_name or not pb_2_grpc_file_name:
                print("Generating client library...")
                lib_generator.generate_client_library(service_api_source)

        if options is None:
            options = dict()
        options["concurrency"] = config.CONCURRENCY
        options["concurrent_calls"] = concurrent_calls

        if payment_strategy is None:
            payment_strategy = payment_strategy_type.value()

        group = self._get_service_group(org_id, service_id, service_metadata, group_name)

        service_stubs = self.get_service_stub(lib_generator)

        pb2_module = self.get_module_by_keyword("pb2.py", lib_generator)
        _service_client = ServiceClient(
            org_id,
            service_id,
            group,
            service_stubs,
            payment_strategy,
            options,
            self.mpe_contract,
            self.account,
            self.w3,
            pb2_module,
            self.payment_channel_provider,
            lib_generator.proto_dir,
            lib_generator.training_added(),
        )
        return _service_client

    def _enhance_service_metadata(self, org_id, service_id):
        service_metadata = self.get_service_metadata(org_id, service_id)
        org_metadata = self.get_organization_metadata(org_id)

        org_group_map = {}
        for group in org_metadata.groups:
            org_group_map[group.group_name] = group

        for group in service_metadata.groups:
            group.payment = org_group_map[group.group_name].payment

        return service_metadata

    def get_service_stub(self, lib_generator: ClientLibGenerator) -> list[ServiceStub]:
        path_to_pb_files = str(lib_generator.proto_dir)
        module_name = self.get_module_by_keyword("pb2_grpc.py", lib_generator)
        sys.path.append(path_to_pb_files)
        try:
            grpc_file = importlib.import_module(module_name)
            properties_and_methods_of_grpc_file = dir(grpc_file)
            service_stubs = []
            for elem in properties_and_methods_of_grpc_file:
                if "Stub" in elem:
                    service_stubs.append(getattr(grpc_file, elem))
            return [ServiceStub(service_stub) for service_stub in service_stubs]

        except Exception as e:
            raise Exception(f"Error importing module: {e}")

    def get_module_by_keyword(self, keyword: str, lib_generator: ClientLibGenerator) -> ModuleName:
        path_to_pb_files = lib_generator.proto_dir
        file_name = find_file_by_keyword(path_to_pb_files, keyword, exclude=["training"])
        module_name = os.path.splitext(file_name)[0]
        return ModuleName(module_name)

    def get_service_metadata(self, org_id, service_id) -> ServiceMetadata:
        service = self.registry_contract.get_service(org_id, service_id)
        return self.storage_provider.fetch_service_metadata(service.metadata_uri)

    def get_organization_metadata(self, org_id: str) -> OrganizationMetadata:
        org = self.registry_contract.get_org(org_id)
        return self.storage_provider.fetch_org_metadata(org.metadata_uri)

    def _get_service_group(
        self, org_id: str, service_id: str, service_metadata: ServiceMetadata, group_name: str
    ) -> Group:
        if len(service_metadata.groups) == 0:
            raise NoGroupsFoundError(org_id, service_id)

        if group_name is None:
            return service_metadata.groups[0]

        for group in service_metadata.groups:
            if group.group_name == group_name:
                return group

        raise GroupNotFoundError(org_id, service_id, group_name)

    def get_organization_list(self) -> list:
        return self.registry_contract.list_orgs()

    def get_services_list(self, org_id: str) -> list:
        return self.registry_contract.list_service_for_org(org_id)

    def publish_service_comprehensively(
        self,
        org_id: str,
        service_id: str,
        metadata: ServiceMetadata,
        proto_dir: Union[str, Path],
        storage_type: StorageType = StorageType.IPFS,
    ) -> bool:
        """
        1. publish .proto files as .tar.gz archive into the storage
        2. add other fields to the service metadata
        3. validate service metadata
        4. publish service metadata into the storage
        5. publish service into Registry contract
        """
        proto_uri = self.storage_provider.publish_proto(proto_dir, storage_type)
        metadata.service_api_source = str(proto_uri)
        metadata.mpe_address = self.mpe_contract.contract.address

        self._check_and_update_service_groups(org_id, metadata.groups)
        metadata_uri = self.storage_provider.publish_service_metadata(metadata, storage_type)

        receipt = self.registry_contract.create_service(
            self.account, org_id, service_id, metadata_uri
        )

        return receipt["status"] != 0

    def update_service(
        self,
        org_id: str,
        service_id: str,
        metadata: ServiceMetadata,
        proto_dir: Union[str, Path, None] = None,
        storage_type: StorageType = StorageType.IPFS,
    ) -> bool:
        if proto_dir is not None:
            proto_uri = self.storage_provider.publish_proto(proto_dir, storage_type)
            metadata.service_api_source = str(proto_uri)

        if not metadata.mpe_address:
            metadata.mpe_address = self.mpe_contract.contract.address

        self._check_and_update_service_groups(org_id, metadata.groups)
        metadata_uri = self.storage_provider.publish_service_metadata(metadata, storage_type)

        receipt = self.registry_contract.update_service_metadata(
            self.account, org_id, service_id, metadata_uri
        )

        return receipt["status"] != 0

    def update_organization(
        self,
        org_id: str,
        organization_metadata: OrganizationMetadata,
        storage_type: StorageType = StorageType.IPFS,
    ) -> bool:
        metadata_uri = self.storage_provider.publish_organization_metadata(
            organization_metadata, storage_type
        )
        receipt = self.registry_contract.update_org_metadata(self.account, org_id, metadata_uri)

        return receipt["status"] != 0

    def _check_and_update_service_groups(
        self, org_id: str, service_groups: list[Group]
    ) -> list[Group]:
        org = self.registry_contract.get_org(org_id)
        org_metadata = self.storage_provider.fetch_org_metadata(org.metadata_uri)
        org_groups_map = {g.group_name: g for g in org_metadata.groups}

        for group in service_groups:
            try:
                group.group_id = org_groups_map[group.group_name].group_id
            except KeyError:
                raise ServiceMetadataMismatchError(
                    "All groups added to the service must also exist in the organization!"
                )

        return service_groups