"""
Operators for the uptrain.core module. 
"""

from __future__ import annotations
import importlib
import types
import typing as t
import typing_extensions as te

from loguru import logger
from pydantic import ConfigDict, BaseModel
import polars as pl

if t.TYPE_CHECKING:
    from uptrain.framework import Settings


__all__ = [
    "TYPE_TABLE_OUTPUT",
    "Operator",
    "OpBaseModel",
    "ColumnOp",
    "TransformOp",
    "PlaceholderOp",
    "register_op",
    "deserialize_operator",
    "get_output_col_name_at",
]

# -----------------------------------------------------------
# Base classes for operators
# -----------------------------------------------------------


def make_module_for_custom_ops():
    """Creates a module to hold custom operators. `inspect.getsource` is
    giving me a lot of grief.
    """
    mod = types.ModuleType("_uptrain_custom_ops")
    preamble = [
        "from uptrain.operators import ColumnOp, TransformOp, register_custom_op",
        "import polars as pl",
    ]
    for line in preamble:
        exec(line, mod.__dict__)
    return mod


class TYPE_TABLE_OUTPUT(te.TypedDict):
    output: pl.DataFrame | None
    extra: te.NotRequired[dict]


class Operator(t.Protocol):
    """
    All operator implementations must implement this protocol.

    An operator is usually defined as a pydantic model with the parameters
    necessary to run it. For ex,
    - params necessary to run an algorithm
    - columns required in the input dataset
    - columns that will be generated as output, etc.


    `Setup`: The `setup` method is called before the operator is run. This is where you
    can do any setup work, like loading models, settin up file addresses, etc.

    `Run`: The `run` method is called to run the operator. It takes 0 or more polars
    dataframes as input, and returns a dictionary with the `output` key set to the
    computed dataframe/None. Any extra information can be put in the `extra` key.
    """

    def dict(self) -> dict:
        """
        Serialize this operator to a json dictionary. The objective is to be able to
        recreate the operator from this dict.

        NOTE: If your operator is a pydantic model, the `to_py_types` method in
        `uptrain.utilities` handles this. Though if you have fields with custom
        non-python types, you MUST override and implement both a `dict` and a`from_dict`
        method.
        TODO: With pydantic v2, nested models work with custom (de)serializers. Use that
        when moving.
        """
        ...

    def setup(self, settings: "Settings") -> "Operator":
        """Setup the operator. This must be called before the operator is run."""
        ...

    run: t.Callable[..., TYPE_TABLE_OUTPUT]  # runs the operator


class OpBaseModel(BaseModel):
    """
    Base class you can use if constructing an operator using a pydantic
    model, to get around some of the sharp edges.
    """

    # TODO[pydantic]: The following keys were removed: `smart_union`, `underscore_attrs_are_private`.
    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information.
    model_config = ConfigDict(extra="allow", protected_namespaces=())


class ColumnOp(OpBaseModel):
    """Represents operations that append columns to the input dataset, and
    return it as is.
    """

    def setup(self, settings: "Settings"):
        raise NotImplementedError

    def run(self, *args: pl.DataFrame) -> TYPE_TABLE_OUTPUT:
        """
        Runs the operator on the input dataset, and returns it with new columns added.

        Args:
            data (pl.DataFrame): Zero or one dataframe as input.

        Returns:
            A dictionary with the `output` key set to the computed Table. Any extra
                information can be put in the `extra` key.
        """
        raise NotImplementedError


class TransformOp(OpBaseModel):
    """Represents operations that transform the input dataset into another.
    For things like filtering, aggregation, etc.
    """

    def setup(self, settings: Settings):
        raise NotImplementedError

    def run(self, *args: pl.DataFrame) -> TYPE_TABLE_OUTPUT:
        """Runs the operator on the input dataset, and returns the transformed dataset.

        Attributes:
            *args (pl.DataFrame): Zero or more dataframes as inputs.

        Returns:
            A dictionary with the `output` key set to the computed dataframe/None. Any extra
                information can be put in the `extra` key.
        """
        raise NotImplementedError


T = t.TypeVar("T")


def register_op(cls: T) -> T:
    """Decorator that marks the class as an Uptrain operator. Useful for (de)serialization."""
    assert isinstance(cls, type), "Only classes can be registered as Uptrain operators"
    assert hasattr(cls, "setup") and hasattr(
        cls, "run"
    ), "All Uptrain operators must define a `setup` and a `run` method."
    op_name = f"{cls.__module__}:{cls.__name__}"
    cls._uptrain_op_name = op_name  # type: ignore
    return cls


def register_custom_op(cls: T) -> T:
    """Decorator that marks the class as a custom Uptrain operator, that is not
    part of the core uptrain package. These are serialized by storing the entire
    source code of the class, and deserialized by exec-ing the source code.

    NOTE: This introduces some restrictions on the operator.
    - This operator is dserialised in an empty namespace, so all imports must
    be done inside the class definition.
    - For typing, use string annotations.
    - Custom operators must be defined in a separate .py file and not in the
    interpreter (like a jupyter notebook).
    """
    cls._uptrain_op_custom = True  # type: ignore
    return register_op(cls)


CUSTOM_OP_MODULE = make_module_for_custom_ops()


def deserialize_operator(data: dict) -> Operator:
    """Deserialize an operator from the serialized dict."""
    op_name = data["op_name"]
    params = data["params"]
    mod_name, cls_name = op_name.split(":")
    try:
        # Check if it is a custom operator, that'd need exec-ing
        if "source" in data:
            exec(data["source"], CUSTOM_OP_MODULE.__dict__)
            klass = getattr(CUSTOM_OP_MODULE, cls_name)
            # to get around serialisation round-trip woes
            klass._uptrain_op_custom_source = data["source"]
            return klass(**params)  # type: ignore
        else:
            ## support for previous operators
            if mod_name.startswith("uptrain.operators"):
                mod_name = "uptrain.operators"
            mod = importlib.import_module(mod_name)
            op = getattr(mod, cls_name)
            # Is it a class implemented by uptrain or a regular pydantic model?
            if hasattr(op, "from_dict"):
                return op.from_dict(params)  # type: ignore
            else:
                return op(**params)  # type: ignore

    except (ModuleNotFoundError, AttributeError) as exc:
        logger.error(
            f"Error when trying to fetch the operator: \n{exc}\n Creating a placeholder op for {op_name}."
        )
        return PlaceholderOp(op_name=op_name, params=data["params"])


# -----------------------------------------------------------
# Utility routines for operators
# -----------------------------------------------------------


def get_output_col_name_at(index: int) -> str:
    return f"_col_{index}"


@register_op
class PlaceholderOp(OpBaseModel):
    """A placeholder operator that does nothing. Used when deserializing an operator,
    that hasn't been registered."""

    op_name: str
    params: dict[str, t.Any]

    def setup(self, settings: "Settings"):
        raise NotImplementedError

    def run(self, *args: pl.DataFrame) -> TYPE_TABLE_OUTPUT:
        raise NotImplementedError(
            "This is only a placeholder since the operator: {self.op_name} is not registered with Uptrain. Import the module where it is defined and try again."
        )
