ResInsight/GrpcInterface/Python/rips/pdmobject.py

# pylint: disable=no-self-use
"""
ResInsight caf::PdmObject connection module
"""

from enum import Enum
from functools import wraps
import grpc
import re
import inspect

import PdmObject_pb2
import PdmObject_pb2_grpc
import Commands_pb2
import Commands_pb2_grpc

from .exception import RipsError

from typing import Any, Callable, TypeVar, Union, List, Optional, Type
from typing_extensions import ParamSpec

# TypeVar for generic return types in PdmObjectBase methods
# Using string bound for forward reference since PdmObjectBase is defined later
PdmObjectT = TypeVar("PdmObjectT", bound="PdmObjectBase")


def camel_to_snake(name: str) -> str:
    s1 = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name)
    return re.sub("([a-z0-9])([A-Z])", r"\1_\2", s1).lower()


def snake_to_camel(name: str) -> str:
    return "".join(word.title() for word in name.split("_"))


F = TypeVar("F", bound=Callable[..., Any])
C = TypeVar("C")


def _is_closed_channel_error(exc: BaseException) -> bool:
    # gRPC raises ValueError("Cannot invoke RPC on closed channel!") when an
    # RPC is attempted after the channel has been closed (e.g. by the
    # heartbeat after detecting a dead server). Translate to RipsError so
    # callers see a typed exception instead of a bare ValueError.
    return isinstance(exc, ValueError) and "closed channel" in str(exc)


def _rips_error_from_rpc_error(exc: grpc.RpcError) -> RipsError:
    """Translate a gRPC failure into a RipsError, adding context for the
    failure modes whose raw gRPC text says nothing useful."""
    details = exc.details() or ""
    message = "%s" % details

    if "Channel closed" in details or exc.code() == grpc.StatusCode.UNAVAILABLE:
        # The connection to ResInsight went away mid-call. This is either a
        # crashed/closed server, or the local channel being closed by the
        # heartbeat. Say so instead of surfacing a bare "Channel closed!".
        message = (
            f"{details}: lost the connection to ResInsight while the call was "
            "running (the server may have crashed or been shut down)"
        )
    elif exc.code() == grpc.StatusCode.RESOURCE_EXHAUSTED:
        message = f"{details}: the gRPC message size limit was exceeded"

    return RipsError(message, code=exc.code(), details=details or None)


def add_method(cls: C) -> Callable[[F], F]:
    def decorator(func: F) -> F:
        def wrapper(*args, **kwargs):
            try:
                return func(*args, **kwargs)
            except grpc.RpcError as e:
                raise _rips_error_from_rpc_error(e) from None
            except ValueError as e:
                if _is_closed_channel_error(e):
                    raise RipsError(
                        "ResInsight gRPC channel is closed "
                        "(server may have crashed or been shut down)"
                    ) from e
                raise

        # Explicitly preserve signature for Sphinx documentation
        wrapper.__name__ = func.__name__
        wrapper.__doc__ = func.__doc__
        wrapper.__signature__ = inspect.signature(func)  # type: ignore
        wrapper.__annotations__ = getattr(func, "__annotations__", {})

        # Add the wrapped function to the class
        setattr(cls, func.__name__, wrapper)
        return func  # returning func means func can still be used normally

    return decorator


P = ParamSpec("P")
T = TypeVar("T")


def add_static_method(cls: C) -> Callable[[Callable[P, T]], Callable[P, T]]:
    def decorator(func: Callable[P, T]) -> Callable[P, T]:
        @wraps(func)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
            return func(*args, **kwargs)

        setattr(cls, func.__name__, wrapper)
        # Note we are not binding func, but wrapper which accepts self but does exactly the same as func
        return func  # returning func means func can still be used normally

    return decorator


class PdmObjectBase:
    """
    The ResInsight base class for the Project Data Model
    """

    __custom_init__ = None

    def _execute_command(self, **command_params) -> Any:
        self.__warnings = []
        response, call = self._commands.Execute.with_call(
            Commands_pb2.CommandParams(**command_params)
        )
        for key, value in call.trailing_metadata():
            value = value.replace(";;", "\n")
            if key == "warning":
                self.__warnings.append(value)

        return response

    def __init__(
        self,
        pb2_object: Optional[PdmObject_pb2.PdmObject],
        channel: Optional[grpc.Channel],
    ) -> None:
        self.__warnings = []
        self.__chunk_size = 8160

        self._channel = channel

        # Create stubs
        if self._channel:
            self._pdm_object_stub = PdmObject_pb2_grpc.PdmObjectServiceStub(
                self._channel
            )
            self._commands = Commands_pb2_grpc.CommandsStub(self._channel)

        if pb2_object is not None:
            # Copy parameters from ResInsight
            assert isinstance(pb2_object, PdmObject_pb2.PdmObject)
            self._pb2_object = pb2_object
            for camel_keyword in self._pb2_object.parameters:
                snake_keyword = camel_to_snake(camel_keyword)
                setattr(self, snake_keyword, self.__get_grpc_value(camel_keyword))
        else:
            # Copy parameters from PdmObject defaults
            self._pb2_object = PdmObject_pb2.PdmObject(
                class_keyword=self.__class__.__name__
            )
            self.__copy_to_pb2()

    def copy_from(self, obj: object) -> None:
        """Copy attribute values from object to self"""
        for attribute in dir(obj):
            if not attribute.startswith("__"):
                value = getattr(obj, attribute)
                # This is crucial to avoid overwriting methods
                if not callable(value):
                    setattr(self, attribute, value)
        if self.__custom_init__ is not None:
            self.__custom_init__(self._pb2_object, self._channel)
        self.update()

    def warnings(self) -> List[str]:
        return self.__warnings

    def has_warnings(self) -> bool:
        return len(self.__warnings) > 0

    def __copy_to_pb2(self) -> None:
        if self._pb2_object is not None:
            for snake_kw in dir(self):
                if not snake_kw.startswith("_"):
                    value = getattr(self, snake_kw)
                    # This is crucial to avoid overwriting methods
                    if not callable(value):
                        camel_kw = snake_to_camel(snake_kw)
                        self.__set_grpc_value(camel_kw, value)

    def pb2_object(self) -> PdmObject_pb2.PdmObject:
        """Private method"""
        return self._pb2_object

    def channel(self) -> grpc.Channel:
        """Private method"""
        return self._channel

    def address(self) -> Any:
        """Get the unique address of the PdmObject

        Returns:
            A 64-bit unsigned integer address
        """

        return self._pb2_object.address

    def set_visible(self, visible: bool) -> None:
        """Set the visibility of the object in the ResInsight project tree"""
        self._pb2_object.visible = visible

    def visible(self) -> bool:
        """Get the visibility of the object in the ResInsight project tree"""
        return bool(self._pb2_object.visible)

    def print_object_info(self) -> None:
        """Print the structure and data content of the PdmObject"""
        print("=========== " + self.__class__.__name__ + " =================")
        print("Object Attributes: ")
        for snake_kw in dir(self):
            if not snake_kw.startswith("_") and not callable(getattr(self, snake_kw)):
                print(
                    "   "
                    + snake_kw
                    + " ["
                    + type(getattr(self, snake_kw)).__name__
                    + "]: "
                    + str(getattr(self, snake_kw))
                )
        print("Object Methods:")
        for snake_kw in dir(self):
            if not snake_kw.startswith("_") and callable(getattr(self, snake_kw)):
                print("   " + snake_kw)

    Value = Union[bool, str, float, int, "ValueArray"]
    ValueArray = List[Value]

    def __convert_from_grpc_value(self, value: str) -> Value:
        if value.lower() == "false":
            return False
        if value.lower() == "true":
            return True
        try:
            int_val = int(value)
            return int_val
        except ValueError:
            try:
                float_val = float(value)
                return float_val
            except ValueError:
                # We may have a string. Remove the outer pair of quotes
                if len(value) >= 2 and value[0] == '"' and value[-1] == '"':
                    value = value[1:-1]
                if self.__islist(value):
                    return self.__makelist(value)
                if self.__istuple(value):
                    return self.__maketuple(value)
                return self.__unescape_string(value)

    def __convert_to_grpc_value(self, value: Any) -> str:
        if isinstance(value, bool):
            if value:
                return "true"
            return "false"
        if isinstance(value, Enum):
            return str(value.value)
        if isinstance(value, PdmObjectBase):
            return value.__class__.__name__ + ":" + str(value.address())
        if isinstance(value, list):
            list_of_values = []
            for val in value:
                list_of_values.append(self.__convert_to_grpc_value(val))
            return "[" + ", ".join(list_of_values) + "]"
        if isinstance(value, tuple):
            list_of_values = []
            for val in value:
                list_of_values.append(self.__convert_to_grpc_value(val))
            return "(" + ", ".join(list_of_values) + ")"
        return str(value)

    def __get_grpc_value(self, camel_keyword: str) -> Value:
        return self.__convert_from_grpc_value(
            self._pb2_object.parameters[camel_keyword]
        )

    def __set_grpc_value(self, camel_keyword: str, value: str) -> None:
        self._pb2_object.parameters[camel_keyword] = self.__convert_to_grpc_value(value)

    def set_value(self, snake_keyword: str, value: object) -> None:
        """Set the value associated with the provided keyword and updates ResInsight

        Arguments:
            keyword(str): A string containing the parameter keyword
            value(varying): A value matching the type of the parameter.
                See keyword documentation and/or print_object_info() to find
                the correct data type.
        """
        setattr(self, snake_keyword, value)
        self.update()

    def __istuple(self, value) -> bool:
        return value.startswith("(") and value.endswith(")")

    def __islist(self, value: str) -> bool:
        return value.startswith("[") and value.endswith("]")

    def __maketuple(self, tuple_string: str) -> Value:
        tuple_string = tuple_string.removeprefix("(")
        tuple_string = tuple_string.removesuffix(")")
        if not tuple_string:
            return ()

        if "), (" in tuple_string:
            subtuples = re.split(r"\), \(", tuple_string)
            return [self.__makelist(subtuple) for subtuple in subtuples]
        else:
            strings = tuple_string.split(", ")
            if len(strings) == 2:
                return (
                    self.__convert_from_grpc_value(strings[0]),
                    self.__convert_from_grpc_value(strings[1]),
                )

        return ()

    def __unescape_string(self, value: str) -> str:
        result = []
        i = 0
        while i < len(value):
            if value[i] == "\\" and i + 1 < len(value):
                next_ch = value[i + 1]
                if next_ch == '"' or next_ch == "\\":
                    result.append(next_ch)
                    i += 2
                    continue
            result.append(value[i])
            i += 1
        return "".join(result)

    def __makelist(self, list_string: str) -> Value:
        list_string = list_string.removeprefix("[")
        list_string = list_string.removesuffix("]")
        if not list_string:
            return []

        # Quote-aware split: track quote state, escape state, and bracket depth
        # so commas inside "..." or [...] are not treated as separators.
        items = []
        current = []
        in_quotes = False
        escape_next = False
        bracket_depth = 0

        for ch in list_string:
            if escape_next:
                current.append(ch)
                escape_next = False
                continue

            if ch == "\\" and in_quotes:
                current.append(ch)
                escape_next = True
                continue

            if ch == '"':
                in_quotes = not in_quotes
                current.append(ch)
                continue

            if not in_quotes:
                if ch == "[":
                    bracket_depth += 1
                elif ch == "]":
                    bracket_depth -= 1
                elif ch == "," and bracket_depth == 0:
                    # Separator: expect ", " so skip the following space
                    items.append("".join(current))
                    current = []
                    continue
                elif ch == " " and not current:
                    # Skip space after comma separator
                    continue

            current.append(ch)

        if current:
            items.append("".join(current))

        values = []
        for item in items:
            values.append(self.__convert_from_grpc_value(item))
        return values

    def __from_pb2_to_resinsight_classes(
        self,
        pb2_object_list: List[PdmObject_pb2.PdmObject],
        super_class_definition: Type[PdmObjectT],
    ) -> List[PdmObjectT]:
        pdm_object_list = []
        from .generated.generated_classes import class_from_keyword

        for pb2_object in pb2_object_list:
            child_class_definition = class_from_keyword(pb2_object.class_keyword)
            if child_class_definition is None:
                child_class_definition = super_class_definition

            assert child_class_definition is not None
            pdm_object = child_class_definition(
                pb2_object=pb2_object, channel=self.channel()
            )

            pdm_object_list.append(pdm_object)
        return pdm_object_list

    def descendants(self, class_definition: Type[PdmObjectT]) -> List[PdmObjectT]:
        """Get a list of all project tree descendants matching the class keyword
        Arguments:
            class_definition[class]: A class definition matching the type of class wanted

        Returns:
            A list of PdmObjects matching the class_definition
        """
        assert inspect.isclass(class_definition)

        class_keyword = class_definition.__name__
        try:
            request = PdmObject_pb2.PdmDescendantObjectRequest(
                object=self._pb2_object, child_keyword=class_keyword
            )
            object_list = self._pdm_object_stub.GetDescendantPdmObjects(request).objects
            return self.__from_pb2_to_resinsight_classes(object_list, class_definition)
        except grpc.RpcError as e:
            if e.code() == grpc.StatusCode.NOT_FOUND:
                return []  # Valid empty result
            raise e

    def children(
        self, child_field: str, class_definition: Type[PdmObjectT]
    ) -> List[PdmObjectT]:
        """Get a list of all direct project tree children inside the provided child_field
        Arguments:
            child_field[str]: A field name
        Returns:
            A list of PdmObjects inside the child_field
        """
        request = PdmObject_pb2.PdmChildObjectRequest(
            object=self._pb2_object, child_field=child_field
        )
        try:
            object_list = self._pdm_object_stub.GetChildPdmObjects(request).objects
            return self.__from_pb2_to_resinsight_classes(object_list, class_definition)
        except grpc.RpcError as e:
            if e.code() == grpc.StatusCode.NOT_FOUND:
                return []
            raise e

    def add_new_object(
        self, class_definition: Type[PdmObjectT], child_field: str = ""
    ) -> Optional[PdmObjectT]:
        """Create and add an object to the specified child field
        Arguments:
            class_definition[class]: Class definition of the object to create
            child_field[str]: The keyword for the field to create a new object in. If empty, the first child array field is used.
        Returns:
            The created PdmObject inside the child_field
        """
        from .generated.generated_classes import class_from_keyword

        assert inspect.isclass(class_definition)

        class_keyword = class_definition.__name__

        request = PdmObject_pb2.CreatePdmChildObjectRequest(
            object=self._pb2_object,
            child_field=child_field,
            class_keyword=class_keyword,
        )
        try:
            pb2_object = self._pdm_object_stub.CreateChildPdmObject(request)
            child_class_definition = class_from_keyword(pb2_object.class_keyword)

            if child_class_definition is None:
                child_class_definition = class_definition

            assert child_class_definition.__name__ == class_definition.__name__
            pdm_object = class_definition(pb2_object=pb2_object, channel=self.channel())

            return pdm_object
        except grpc.RpcError as e:
            if e.code() == grpc.StatusCode.NOT_FOUND:
                return None
            raise e

    def ancestor(self, class_definition: Type[PdmObjectT]) -> Optional[PdmObjectT]:
        """Find the first ancestor that matches the provided class_keyword
        Arguments:
            class_definition[class]: A class definition matching the type of class wanted
        """
        assert inspect.isclass(class_definition)

        class_keyword = class_definition.__name__

        request = PdmObject_pb2.PdmParentObjectRequest(
            object=self._pb2_object, parent_keyword=class_keyword
        )
        try:
            pb2_object = self._pdm_object_stub.GetAncestorPdmObject(request)
            pdm_object = class_definition(pb2_object=pb2_object, channel=self.channel())
            return pdm_object
        except grpc.RpcError as e:
            if e.code() == grpc.StatusCode.NOT_FOUND:
                return None
            raise e

    def _call_get_method_async(self, method_name: str):
        request = PdmObject_pb2.PdmObjectGetterRequest(
            object=self._pb2_object, method=method_name
        )
        for chunk in self._pdm_object_stub.CallPdmObjectGetter(request):
            yield chunk

    def _call_get_method(self, method_name: str):
        all_values = []
        generator = self._call_get_method_async(method_name)
        for chunk in generator:
            data = getattr(chunk, chunk.WhichOneof("data"))
            for value in data.data:
                all_values.append(value)
        return all_values

    def __generate_set_method_chunks(self, array, method_request):
        index = -1

        while index < len(array):
            chunk = PdmObject_pb2.PdmObjectSetterChunk()
            if index == -1:
                chunk.set_request.CopyFrom(
                    PdmObject_pb2.PdmObjectSetterRequest(
                        request=method_request, data_count=len(array)
                    )
                )
                index += 1
            else:
                actual_chunk_size = min(len(array) - index + 1, self.__chunk_size)
                if isinstance(array[0], float):
                    chunk.CopyFrom(
                        PdmObject_pb2.PdmObjectSetterChunk(
                            doubles=PdmObject_pb2.DoubleArray(
                                data=array[index : index + actual_chunk_size]
                            )
                        )
                    )
                elif isinstance(array[0], int):
                    chunk.CopyFrom(
                        PdmObject_pb2.PdmObjectSetterChunk(
                            ints=PdmObject_pb2.IntArray(
                                data=array[index : index + actual_chunk_size]
                            )
                        )
                    )
                elif isinstance(array[0], str):
                    chunk.CopyFrom(
                        PdmObject_pb2.PdmObjectSetterChunk(
                            strings=PdmObject_pb2.StringArray(
                                data=array[index : index + actual_chunk_size]
                            )
                        )
                    )
                else:
                    raise Exception("Wrong data type for set method")
                index += actual_chunk_size
            yield chunk
        # Final empty message to signal completion
        chunk = PdmObject_pb2.PdmObjectSetterChunk()
        yield chunk

    def _call_set_method(self, method_name: str, values) -> None:
        method_request = PdmObject_pb2.PdmObjectGetterRequest(
            object=self._pb2_object, method=method_name
        )
        request_iterator = self.__generate_set_method_chunks(values, method_request)
        reply = self._pdm_object_stub.CallPdmObjectSetter(request_iterator)
        if reply.accepted_value_count < len(values):
            raise IndexError

    def _call_pdm_method_void(self, method_name: str, **kwargs: Any) -> None:
        pb2_params = PdmObject_pb2.PdmObject(class_keyword=method_name)
        for key, value in kwargs.items():
            pb2_params.parameters[snake_to_camel(key)] = self.__convert_to_grpc_value(
                value
            )
        request = PdmObject_pb2.PdmObjectMethodRequest(
            object=self._pb2_object, method=method_name, params=pb2_params
        )

        try:
            self._pdm_object_stub.CallPdmObjectMethod(request)
        except grpc.RpcError as exc:
            raise _rips_error_from_rpc_error(exc) from None
        except ValueError as exc:
            if _is_closed_channel_error(exc):
                raise RipsError(
                    "ResInsight gRPC channel is closed "
                    "(server may have crashed or been shut down)"
                ) from exc
            raise

    def _call_pdm_method_return_value(
        self, method_name: str, class_definition: Type[PdmObjectT], **kwargs: Any
    ) -> PdmObjectT:
        pb2_params = PdmObject_pb2.PdmObject(class_keyword=method_name)
        for key, value in kwargs.items():
            pb2_params.parameters[snake_to_camel(key)] = self.__convert_to_grpc_value(
                value
            )
        request = PdmObject_pb2.PdmObjectMethodRequest(
            object=self._pb2_object, method=method_name, params=pb2_params
        )

        try:
            pb2_object = self._pdm_object_stub.CallPdmObjectMethod(request)

            # Prefer the actual class the server returned over the statically declared one
            # so that generic methods (e.g. AddFolder registered on a shared base) yield
            # the most-derived Python proxy and inherited accessors work as expected.
            from .generated.generated_classes import class_from_keyword

            actual_class = class_from_keyword(pb2_object.class_keyword)
            if actual_class is not None and issubclass(actual_class, class_definition):
                return actual_class(pb2_object=pb2_object, channel=self.channel())

            return class_definition(pb2_object=pb2_object, channel=self.channel())
        except grpc.RpcError as exc:
            raise _rips_error_from_rpc_error(exc) from None
        except ValueError as exc:
            if _is_closed_channel_error(exc):
                raise RipsError(
                    "ResInsight gRPC channel is closed "
                    "(server may have crashed or been shut down)"
                ) from exc
            raise

    def _call_pdm_method_return_optional_value(
        self, method_name: str, class_definition: Type[PdmObjectT], **kwargs: Any
    ) -> Optional[PdmObjectT]:
        pb2_params = PdmObject_pb2.PdmObject(class_keyword=method_name)
        for key, value in kwargs.items():
            pb2_params.parameters[snake_to_camel(key)] = self.__convert_to_grpc_value(
                value
            )
        request = PdmObject_pb2.PdmObjectMethodRequest(
            object=self._pb2_object, method=method_name, params=pb2_params
        )

        try:
            pb2_object = self._pdm_object_stub.CallPdmObjectMethod(request)

            from .generated.generated_classes import class_from_keyword

            child_class_definition = class_from_keyword(pb2_object.class_keyword)
            if child_class_definition is None:
                return None

            assert issubclass(child_class_definition, class_definition)
            pdm_object = child_class_definition(
                pb2_object=pb2_object, channel=self.channel()
            )
            return pdm_object

        except grpc.RpcError as exc:
            raise _rips_error_from_rpc_error(exc) from None
        except ValueError as exc:
            if _is_closed_channel_error(exc):
                raise RipsError(
                    "ResInsight gRPC channel is closed "
                    "(server may have crashed or been shut down)"
                ) from exc
            raise

    def update(self) -> None:
        """Sync all fields from the Python Object to ResInsight

        If validation fails on the server, all local changes are rolled back
        to the last known valid state, and the exception is re-raised.
        """
        if self._pdm_object_stub is None:
            raise Exception(
                "Object is not connected to GRPC service so cannot update ResInsight"
            )

        # Store a snapshot of the last known valid protobuf state
        pb2_snapshot = PdmObject_pb2.PdmObject()
        pb2_snapshot.CopyFrom(self._pb2_object)

        try:
            self.__copy_to_pb2()
            self._pdm_object_stub.UpdateExistingPdmObject(self._pb2_object)
        except Exception:
            # Rollback: restore protobuf object to previous state
            self._pb2_object.CopyFrom(pb2_snapshot)
            # Restore Python attributes from the protobuf snapshot
            for camel_keyword in pb2_snapshot.parameters:
                snake_keyword = camel_to_snake(camel_keyword)
                value = self.__convert_from_grpc_value(
                    pb2_snapshot.parameters[camel_keyword]
                )
                setattr(self, snake_keyword, value)
            # Re-raise the exception so the caller knows update failed
            raise

    def delete(self) -> None:
        """Delete object in ResInsight"""
        self.__copy_to_pb2()
        if self._pdm_object_stub is not None:
            self._pdm_object_stub.DeleteExistingPdmObject(self._pb2_object)
        else:
            raise Exception(
                "Object is not connected to GRPC service so cannot update ResInsight"
            )