#! /usr/bin/env python3

"""
TODO: szengine_grpc.py
"""

# pylint: disable=E1101,C0302

import json
from types import TracebackType
from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Union

import grpc
from senzing import SzEngine as SzEngineAbstract
from senzing import SzEngineFlags

from .pb2_grpc import szengine_pb2, szengine_pb2_grpc
from .szhelpers import as_str, new_exception

# Metadata

__all__ = ["SzEngine"]
__version__ = "0.0.1"  # See https://www.python.org/dev/peps/pep-0396/
__date__ = "2023-11-27"
__updated__ = "2024-07-30"

SENZING_PRODUCT_ID = (
    "5053"  # See https://github.com/senzing-garage/knowledge-base/blob/main/lists/senzing-component-ids.md
)

# -----------------------------------------------------------------------------
# SzEngine class
# -----------------------------------------------------------------------------


class SzEngine(SzEngineAbstract):
    """
    Sz engine module access library over gRPC.
    """

    # -------------------------------------------------------------------------
    # Python dunder/magic methods
    # -------------------------------------------------------------------------

    def __init__(
        self,
        grpc_channel: grpc.Channel,
    ) -> None:
        """
        Constructor

        For return value of -> None, see https://peps.python.org/pep-0484/#the-meaning-of-annotations
        """

        self.channel = grpc_channel
        self.stub = szengine_pb2_grpc.SzEngineStub(self.channel)
        self.noop = ""

    def __enter__(
        self,
    ) -> Any:  # TODO: Replace "Any" with "Self" once python 3.11 is lowest supported python version.
        """Context Manager method."""
        return self

    def __exit__(
        self,
        exc_type: Union[Type[BaseException], None],
        exc_val: Union[BaseException, None],
        exc_tb: Union[TracebackType, None],
    ) -> None:
        """Context Manager method."""

    # -------------------------------------------------------------------------
    # SzEngine methods
    # -------------------------------------------------------------------------

    def add_record(
        self,
        data_source_code: str,
        record_id: str,
        record_definition: str,
        flags: int = 0,
    ) -> str:
        try:
            request = szengine_pb2.AddRecordRequest(  # type: ignore[unused-ignore]
                dataSourceCode=as_str(data_source_code),
                recordId=as_str(record_id),
                recordDefinition=as_str(record_definition),
                flags=flags,
            )
            response = self.stub.AddRecord(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def close_export(self, export_handle: int) -> None:
        try:
            request = szengine_pb2.CloseExportRequest(  # type: ignore[unused-ignore]
                exportHandle=export_handle,
            )
            self.stub.CloseExport(request)
        except Exception as err:
            raise new_exception(err) from err

    def count_redo_records(self) -> int:
        try:
            request = szengine_pb2.CountRedoRecordsRequest()  # type: ignore[unused-ignore]
            response = self.stub.CountRedoRecords(request)
            return int(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def delete_record(
        self,
        data_source_code: str,
        record_id: str,
        flags: int = 0,
    ) -> str:
        try:
            request = szengine_pb2.DeleteRecordRequest(  # type: ignore[unused-ignore]
                dataSourceCode=as_str(data_source_code),
                recordId=as_str(record_id),
                flags=flags,
            )
            response = self.stub.DeleteRecord(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def _destroy(self) -> None:
        """Null function in the sz-sdk-python-grpc implementation."""

    def export_csv_entity_report(
        self,
        csv_column_list: str,
        flags: int = SzEngineFlags.SZ_EXPORT_DEFAULT_FLAGS,
    ) -> int:
        try:
            request = szengine_pb2.ExportCsvEntityReportRequest(  # type: ignore[unused-ignore]
                csvColumnList=as_str(csv_column_list),
                flags=flags,
            )
            response = self.stub.ExportCsvEntityReport(request)
            return int(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def export_csv_entity_report_iterator(
        self,
        csv_column_list: str,
        flags: int = SzEngineFlags.SZ_EXPORT_DEFAULT_FLAGS,
    ) -> Iterable[str]:
        """TODO: Add method docstring to export_csv_entity_report_iterator."""
        try:
            request = szengine_pb2.StreamExportCsvEntityReportRequest(  # type: ignore[unused-ignore]
                csvColumnList=as_str(csv_column_list), flags=flags
            )
            for item in self.stub.StreamExportCsvEntityReport(request):
                if item.result:
                    yield item.result
        except Exception as err:
            raise new_exception(err) from err

    def export_json_entity_report(self, flags: int = SzEngineFlags.SZ_EXPORT_DEFAULT_FLAGS) -> int:
        try:
            request = szengine_pb2.ExportJsonEntityReportRequest(  # type: ignore[unused-ignore]
                flags=flags,
            )
            response = self.stub.ExportJsonEntityReport(request)
            return int(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def export_json_entity_report_iterator(
        self,
        flags: int = SzEngineFlags.SZ_EXPORT_DEFAULT_FLAGS,
    ) -> Iterable[str]:
        """TODO: Add method docstring to export_json_entity_report_iterator."""
        try:
            request = szengine_pb2.StreamExportJsonEntityReportRequest(flags=flags)  # type: ignore[unused-ignore]
            for item in self.stub.StreamExportJsonEntityReport(request):
                if item.result:
                    yield item.result
        except Exception as err:
            raise new_exception(err) from err

    def fetch_next(self, export_handle: int) -> str:
        try:
            request = szengine_pb2.FetchNextRequest(  # type: ignore[unused-ignore]
                exportHandle=export_handle,
            )
            response = self.stub.FetchNext(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def find_interesting_entities_by_entity_id(self, entity_id: int, flags: int = 0) -> str:
        try:
            request = szengine_pb2.FindInterestingEntitiesByEntityIdRequest(  # type: ignore[unused-ignore]
                entityId=entity_id,
                flags=flags,
            )
            response = self.stub.FindInterestingEntitiesByEntityId(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def find_interesting_entities_by_record_id(self, data_source_code: str, record_id: str, flags: int = 0) -> str:
        try:
            request = szengine_pb2.FindInterestingEntitiesByRecordIdRequest(  # type: ignore[unused-ignore]
                dataSourceCode=as_str(data_source_code),
                recordId=as_str(record_id),
                flags=flags,
            )
            response = self.stub.FindInterestingEntitiesByRecordId(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def find_network_by_entity_id(
        self,
        entity_ids: List[int],
        max_degrees: int,
        build_out_degrees: int,
        build_out_max_entities: int,
        flags: int = SzEngineFlags.SZ_FIND_NETWORK_DEFAULT_FLAGS,
    ) -> str:
        try:
            request = szengine_pb2.FindNetworkByEntityIdRequest(  # type: ignore[unused-ignore]
                entityIds=entity_ids_json(entity_ids),
                maxDegrees=max_degrees,
                buildOutDegrees=build_out_degrees,
                buildOutMaxEntities=build_out_max_entities,
                flags=flags,
            )
            response = self.stub.FindNetworkByEntityId(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def find_network_by_record_id(
        self,
        record_keys: List[Tuple[str, str]],
        max_degrees: int,
        build_out_degrees: int,
        build_out_max_entities: int,
        flags: int = SzEngineFlags.SZ_FIND_NETWORK_DEFAULT_FLAGS,
    ) -> str:
        try:
            request = szengine_pb2.FindNetworkByRecordIdRequest(  # type: ignore[unused-ignore]
                recordKeys=record_keys_json(record_keys),
                maxDegrees=max_degrees,
                buildOutDegrees=build_out_degrees,
                buildOutMaxEntities=build_out_max_entities,
                flags=flags,
            )
            response = self.stub.FindNetworkByRecordId(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def find_path_by_entity_id(
        self,
        start_entity_id: int,
        end_entity_id: int,
        max_degrees: int,
        # TODO Should accept both entity and record IDs in V4, test
        avoid_entity_ids: Optional[List[int]] = None,
        required_data_sources: Optional[List[str]] = None,
        flags: int = SzEngineFlags.SZ_FIND_PATH_DEFAULT_FLAGS,
    ) -> str:
        try:
            request = szengine_pb2.FindPathByEntityIdRequest(  # type: ignore[unused-ignore]
                startEntityId=start_entity_id,
                endEntityId=end_entity_id,
                maxDegrees=max_degrees,
                avoidEntityIds=avoid_entity_ids_json(avoid_entity_ids),
                requiredDataSources=required_data_sources_json(required_data_sources),
                flags=flags,
            )
            response = self.stub.FindPathByEntityId(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def find_path_by_record_id(
        self,
        start_data_source_code: str,
        start_record_id: str,
        end_data_source_code: str,
        end_record_id: str,
        max_degrees: int,
        avoid_record_keys: Optional[List[Tuple[str, str]]] = None,
        required_data_sources: Optional[List[str]] = None,
        flags: int = SzEngineFlags.SZ_FIND_PATH_DEFAULT_FLAGS,
    ) -> str:
        try:
            request = szengine_pb2.FindPathByRecordIdRequest(  # type: ignore[unused-ignore]
                startDataSourceCode=as_str(start_data_source_code),
                startRecordId=as_str(start_record_id),
                endDataSourceCode=as_str(end_data_source_code),
                endRecordId=as_str(end_record_id),
                maxDegrees=max_degrees,
                avoidRecordKeys=avoid_record_keys_json(avoid_record_keys),
                requiredDataSources=required_data_sources_json(required_data_sources),
                flags=flags,
            )
            response = self.stub.FindPathByRecordId(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def get_active_config_id(self) -> int:
        try:
            request = szengine_pb2.GetActiveConfigIdRequest()  # type: ignore[unused-ignore]
            response = self.stub.GetActiveConfigId(request)
            return int(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def get_entity_by_entity_id(
        self,
        entity_id: int,
        flags: int = SzEngineFlags.SZ_ENTITY_DEFAULT_FLAGS,
    ) -> str:
        try:
            request = szengine_pb2.GetEntityByEntityIdRequest(  # type: ignore[unused-ignore]
                entityId=entity_id,
                flags=flags,
            )
            response = self.stub.GetEntityByEntityId(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def get_entity_by_record_id(
        self,
        data_source_code: str,
        record_id: str,
        flags: int = SzEngineFlags.SZ_ENTITY_DEFAULT_FLAGS,
    ) -> str:
        try:
            request = szengine_pb2.GetEntityByRecordIdRequest(  # type: ignore[unused-ignore]
                dataSourceCode=as_str(data_source_code),
                recordId=as_str(record_id),
                flags=flags,
            )
            response = self.stub.GetEntityByRecordId(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def get_record(
        self,
        data_source_code: str,
        record_id: str,
        flags: int = SzEngineFlags.SZ_RECORD_DEFAULT_FLAGS,
    ) -> str:
        try:
            request = szengine_pb2.GetRecordRequest(  # type: ignore[unused-ignore]
                dataSourceCode=as_str(data_source_code),
                recordId=as_str(record_id),
                flags=flags,
            )
            response = self.stub.GetRecord(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def get_redo_record(self) -> str:
        try:
            request = szengine_pb2.GetRedoRecordRequest()  # type: ignore[unused-ignore]
            response = self.stub.GetRedoRecord(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def get_stats(self) -> str:
        try:
            request = szengine_pb2.GetStatsRequest()  # type: ignore[unused-ignore]
            response = self.stub.GetStats(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def get_virtual_entity_by_record_id(
        self,
        record_keys: List[Tuple[str, str]],
        flags: int = SzEngineFlags.SZ_VIRTUAL_ENTITY_DEFAULT_FLAGS,
    ) -> str:
        try:
            request = szengine_pb2.GetVirtualEntityByRecordIdRequest(  # type: ignore[unused-ignore]
                recordKeys=record_keys_json(record_keys),
                flags=flags,
            )
            response = self.stub.GetVirtualEntityByRecordId(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def how_entity_by_entity_id(
        self,
        entity_id: int,
        flags: int = SzEngineFlags.SZ_HOW_ENTITY_DEFAULT_FLAGS,
    ) -> str:
        try:
            request = szengine_pb2.HowEntityByEntityIdRequest(  # type: ignore[unused-ignore]
                entityId=entity_id,
                flags=flags,
            )
            response = self.stub.HowEntityByEntityId(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def _initialize(
        self,
        instance_name: str,
        settings: Union[str, Dict[Any, Any]],
        config_id: Optional[int] = None,
        verbose_logging: int = 0,
    ) -> None:
        """Null function in the sz-sdk-python-grpc implementation."""
        _ = instance_name
        _ = settings
        _ = config_id
        _ = verbose_logging

    def preprocess_record(
        self,
        record_definition: str,
        flags: int = SzEngineFlags.SZ_RECORD_DEFAULT_FLAGS,
    ) -> str:
        try:
            request = szengine_pb2.PreprocessRecordRequest(  # type: ignore[unused-ignore]
                recordDefinition=as_str(record_definition),
                flags=flags,
            )
            response = self.stub.PreprocessRecord(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def prime_engine(self) -> None:
        """Null function in the sz-sdk-python-grpc implementation."""

    def process_redo_record(self, redo_record: str, flags: int = 0) -> str:
        try:
            request = szengine_pb2.ProcessRedoRecordRequest(  # type: ignore[unused-ignore]
                redoRecord=as_str(redo_record),
                flags=flags,
            )
            response = self.stub.ProcessRedoRecord(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def reevaluate_entity(self, entity_id: int, flags: int = 0) -> str:
        try:
            request = szengine_pb2.ReevaluateEntityRequest(  # type: ignore[unused-ignore]
                entityId=entity_id,
                flags=flags,
            )
            response = self.stub.ReevaluateEntity(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def reevaluate_record(self, data_source_code: str, record_id: str, flags: int = 0) -> str:
        try:
            request = szengine_pb2.ReevaluateRecordRequest(  # type: ignore[unused-ignore]
                dataSourceCode=as_str(data_source_code),
                recordId=as_str(record_id),
                flags=flags,
            )
            response = self.stub.ReevaluateRecord(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def _reinitialize(self, config_id: int) -> None:
        try:
            request = szengine_pb2.ReinitializeRequest(configId=config_id)  # type: ignore[unused-ignore]
            self.stub.Reinitialize(request)
        except Exception as err:
            raise new_exception(err) from err

    def search_by_attributes(
        self,
        attributes: str,
        flags: int = SzEngineFlags.SZ_SEARCH_BY_ATTRIBUTES_DEFAULT_FLAGS,
        search_profile: str = "",
    ) -> str:
        try:
            request = szengine_pb2.SearchByAttributesRequest(  # type: ignore[unused-ignore]
                attributes=as_str(attributes),
                searchProfile=as_str(search_profile),
                flags=flags,
            )
            response = self.stub.SearchByAttributes(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def why_entities(
        self,
        entity_id_1: int,
        entity_id_2: int,
        flags: int = SzEngineFlags.SZ_WHY_ENTITIES_DEFAULT_FLAGS,
    ) -> str:
        try:
            request = szengine_pb2.WhyEntitiesRequest(  # type: ignore[unused-ignore]
                entityId1=entity_id_1,
                entityId2=entity_id_2,
                flags=flags,
            )
            response = self.stub.WhyEntities(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def why_record_in_entity(
        self,
        data_source_code: str,
        record_id: str,
        flags: int = SzEngineFlags.SZ_WHY_RECORDS_DEFAULT_FLAGS,
    ) -> str:
        # TODO: Implement after V3 is published.
        try:
            request = szengine_pb2.WhyRecordInEntityRequest(  # type: ignore[unused-ignore]
                dataSourceCode=as_str(data_source_code),
                recordId=as_str(record_id),
                flags=flags,
            )
            response = self.stub.WhyRecordInEntity(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err

    def why_records(
        self,
        data_source_code_1: str,
        record_id_1: str,
        data_source_code_2: str,
        record_id_2: str,
        flags: int = SzEngineFlags.SZ_WHY_RECORDS_DEFAULT_FLAGS,
    ) -> str:
        try:
            request = szengine_pb2.WhyRecordsRequest(  # type: ignore[unused-ignore]
                dataSourceCode1=as_str(data_source_code_1),
                recordId1=as_str(record_id_1),
                dataSourceCode2=as_str(data_source_code_2),
                recordId2=as_str(record_id_2),
                flags=flags,
            )
            response = self.stub.WhyRecords(request)
            return str(response.result)
        except Exception as err:
            raise new_exception(err) from err


# -----------------------------------------------------------------------------
# Helper functions
# -----------------------------------------------------------------------------


def entity_ids_json(entity_ids: List[int]) -> str:
    entity_list = []
    for entity_id in entity_ids:
        entity_list.append({"ENTITY_ID": entity_id})
    return json.dumps({"ENTITIES": entity_list})


def record_keys_json(record_keys: List[Tuple[str, str]]) -> str:
    record_key_list = []
    for record_key in record_keys:
        record_key_list.append({"DATA_SOURCE": record_key[0], "RECORD_ID": record_key[1]})
    return json.dumps({"RECORDS": record_key_list})


def avoid_entity_ids_json(avoid_entity_ids: Optional[List[int]] = None) -> str:
    result = ""
    if avoid_entity_ids:
        avoid_entity_id_list = []
        for avoid_entity_id in avoid_entity_ids:
            avoid_entity_id_list.append({"ENTITY_ID": avoid_entity_id})
        result = json.dumps({"ENTITIES": avoid_entity_id_list})
    return result


def avoid_record_keys_json(avoid_record_keys: Optional[List[Tuple[str, str]]] = None) -> str:
    result = ""
    if avoid_record_keys:
        avoid_record_keys_list = []
        for avoid_record_key in avoid_record_keys:
            avoid_record_keys_list.append(
                {
                    "DATA_SOURCE": avoid_record_key[0],
                    "RECORD_ID": avoid_record_key[1],
                }
            )
        result = json.dumps({"RECORDS": avoid_record_keys_list})
    return result


def required_data_sources_json(
    required_data_sources: Optional[List[str]] = None,
) -> str:
    result = ""
    if required_data_sources:
        required_data_sources_list = []
        for required_data_source in required_data_sources:
            required_data_sources_list.append(required_data_source)
        result = json.dumps({"DATA_SOURCES": required_data_sources_list})
    return result
