# Copyright 2023 Iguazio
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import random
from datetime import timedelta
from typing import Any, Optional, Union

import numpy as np
import storey

import mlrun
import mlrun.artifacts
import mlrun.common.schemas.model_monitoring as mm_schemas
import mlrun.serving
from mlrun.common.schemas import MonitoringData
from mlrun.utils import get_data_from_path, logger


class MonitoringPreProcessor(storey.MapClass):
    """preprocess step, reconstructs the serving output event body to StreamProcessingEvent schema"""

    def __init__(
        self,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.server: mlrun.serving.GraphServer = (
            getattr(self.context, "server", None) if self.context else None
        )

    def reconstruct_request_resp_fields(
        self, event, model: str, model_monitoring_data: dict
    ) -> tuple[dict[str, Any], dict[str, Any]]:
        result_path = model_monitoring_data.get(MonitoringData.RESULT_PATH)
        input_path = model_monitoring_data.get(MonitoringData.INPUT_PATH)

        result = get_data_from_path(result_path, event.body.get(model, event.body))
        output_schema = model_monitoring_data.get(MonitoringData.OUTPUTS)
        input_schema = model_monitoring_data.get(MonitoringData.INPUTS)
        logger.debug("output schema retrieved", output_schema=output_schema)
        if isinstance(result, dict):
            # transpose by key the outputs:
            outputs = self.transpose_by_key(result, output_schema)
            if not output_schema:
                logger.warn(
                    "Output schema was not provided using Project:log_model or by ModelRunnerStep:add_model order "
                    "may not preserved"
                )
        else:
            outputs = result

        event_inputs = event._metadata.get("inputs", {})
        event_inputs = get_data_from_path(input_path, event_inputs)
        if isinstance(event_inputs, dict):
            # transpose by key the inputs:
            inputs = self.transpose_by_key(event_inputs, input_schema)
            if not input_schema:
                logger.warn(
                    "Input schema was not provided using by ModelRunnerStep:add_model, order "
                    "may not preserved"
                )
        else:
            inputs = event_inputs

        if outputs and isinstance(outputs[0], list):
            if output_schema and len(output_schema) != len(outputs[0]):
                logger.info(
                    "The number of outputs returned by the model does not match the number of outputs "
                    "specified in the model endpoint.",
                    model_endpoint=model,
                    output_len=len(outputs[0]),
                    schema_len=len(output_schema),
                )
        elif outputs:
            if output_schema and len(output_schema) != 1:
                logger.info(
                    "The number of outputs returned by the model does not match the number of outputs "
                    "specified in the model endpoint.",
                    model_endpoint=model,
                    output_len=len(outputs),
                    schema_len=len(output_schema),
                )
            if len(inputs) != len(outputs):
                logger.warn(
                    "outputs and inputs are not in the same length check 'input_path' and "
                    "'output_path' was specified if needed"
                )
        request = {"inputs": inputs, "id": getattr(event, "id", None)}
        resp = {"outputs": outputs}

        return request, resp

    @staticmethod
    def transpose_by_key(
        data: dict, schema: Optional[Union[str, list[str]]] = None
    ) -> Union[list[Any], list[list[Any]]]:
        """
        Transpose values from a dictionary by keys.

        Given a dictionary and an optional schema (a key or list of keys), this function:
        - Extracts the values for the specified keys (or all keys if no schema is provided).
        - Ensures the data is represented as a list of rows, then transposes it (i.e., switches rows to columns).
        - Handles edge cases:
            * If a single scalar or single-element list is provided, returns a flat list.
            * If a single key is provided (as a string or a list with one element), handles it properly.
            * If only one row with len of one remains after transposition, unwraps it to avoid nested list-of-one.

        Example::

            transpose_by_key({"a": 1})
            # returns: [1]

            transpose_by_key({"a": [1, 2]})
            # returns: [1 ,2]

            transpose_by_key({"a": [1, 2], "b": [3, 4]})
            # returns: [[1, 3], [2, 4]]

        :param data:     Dictionary with values that are either scalars or lists.
        :param schema:   Optional key or list of keys to extract. If not provided, all keys are used.
                         Can be a string (single key) or a list of strings.

        :return:         Transposed values:
                         * If result is a single column or row, returns a flat list.
                         * If result is a matrix, returns a list of lists.

        :raises ValueError: If the values include a mix of scalars and lists, or if the list lengths do not match.
        """

        # Normalize schema to list
        if not schema:
            keys = list(data.keys())
        elif isinstance(schema, str):
            keys = [schema]
        else:
            keys = schema

        values = [data[key] for key in keys if key in data]
        if len(values) != len(keys):
            raise mlrun.MLRunInvalidArgumentError(
                f"Schema keys {keys} do not match the data keys {list(data.keys())}."
            )

        # Detect if all are scalars ie: int,float,str
        all_scalars = all(not isinstance(v, (list, tuple, np.ndarray)) for v in values)
        all_lists = all(isinstance(v, (list, tuple, np.ndarray)) for v in values)

        if not (all_scalars or all_lists):
            raise ValueError(
                "All values must be either scalars or lists of equal length."
            )

        if all_scalars:
            transposed = np.array([values], dtype=object)
        elif all_lists and len(keys) > 1:
            arrays = [np.array(v, dtype=object) for v in values]
            mat = np.stack(arrays, axis=0)
            transposed = mat.T
        else:
            return values[0]

        if transposed.shape[1] == 1 and transposed.shape[0] == 1:
            # Transform [[0]] -> [0]:
            return transposed[:, 0].tolist()
        return transposed.tolist()

    def do(self, event):
        monitoring_event_list = []
        model_runner_name = event._metadata.get("model_runner_name", "")
        step = self.server.graph.steps[model_runner_name] if self.server else None
        if not step or not hasattr(step, "monitoring_data"):
            raise mlrun.errors.MLRunRuntimeError(
                f"ModelRunnerStep name {model_runner_name} is not found in the graph or does not have monitoring data"
            )
        monitoring_data = step.monitoring_data
        logger.debug(
            "monitoring preprocessor started",
            event=event,
            monitoring_data=monitoring_data,
            metadata=event._metadata,
        )
        if len(monitoring_data) > 1:
            for model in event.body.keys():
                if model in monitoring_data:
                    request, resp = self.reconstruct_request_resp_fields(
                        event, model, monitoring_data[model]
                    )
                    if hasattr(event, "_original_timestamp"):
                        when = event._original_timestamp
                    else:
                        when = event._metadata.get(model, {}).get(
                            mm_schemas.StreamProcessingEvent.WHEN
                        )
                    monitoring_event_list.append(
                        {
                            mm_schemas.StreamProcessingEvent.MODEL: model,
                            mm_schemas.StreamProcessingEvent.MODEL_CLASS: monitoring_data[
                                model
                            ].get(mm_schemas.StreamProcessingEvent.MODEL_CLASS),
                            mm_schemas.StreamProcessingEvent.MICROSEC: event._metadata.get(
                                model, {}
                            ).get(mm_schemas.StreamProcessingEvent.MICROSEC),
                            mm_schemas.StreamProcessingEvent.WHEN: when,
                            mm_schemas.StreamProcessingEvent.ENDPOINT_ID: monitoring_data[
                                model
                            ].get(
                                mlrun.common.schemas.MonitoringData.MODEL_ENDPOINT_UID
                            ),
                            mm_schemas.StreamProcessingEvent.LABELS: monitoring_data[
                                model
                            ].get(mlrun.common.schemas.MonitoringData.OUTPUTS),
                            mm_schemas.StreamProcessingEvent.FUNCTION_URI: self.server.function_uri
                            if self.server
                            else None,
                            mm_schemas.StreamProcessingEvent.REQUEST: request,
                            mm_schemas.StreamProcessingEvent.RESPONSE: resp,
                            mm_schemas.StreamProcessingEvent.ERROR: event.body[model][
                                mm_schemas.StreamProcessingEvent.ERROR
                            ]
                            if mm_schemas.StreamProcessingEvent.ERROR
                            in event.body[model]
                            else None,
                            mm_schemas.StreamProcessingEvent.METRICS: event.body[model][
                                mm_schemas.StreamProcessingEvent.METRICS
                            ]
                            if mm_schemas.StreamProcessingEvent.METRICS
                            in event.body[model]
                            else None,
                        }
                    )
        elif monitoring_data:
            model = list(monitoring_data.keys())[0]
            request, resp = self.reconstruct_request_resp_fields(
                event, model, monitoring_data[model]
            )
            if hasattr(event, "_original_timestamp"):
                when = event._original_timestamp
            else:
                when = event._metadata.get(mm_schemas.StreamProcessingEvent.WHEN)
            monitoring_event_list.append(
                {
                    mm_schemas.StreamProcessingEvent.MODEL: model,
                    mm_schemas.StreamProcessingEvent.MODEL_CLASS: monitoring_data[
                        model
                    ].get(mm_schemas.StreamProcessingEvent.MODEL_CLASS),
                    mm_schemas.StreamProcessingEvent.MICROSEC: event._metadata.get(
                        mm_schemas.StreamProcessingEvent.MICROSEC
                    ),
                    mm_schemas.StreamProcessingEvent.WHEN: when,
                    mm_schemas.StreamProcessingEvent.ENDPOINT_ID: monitoring_data[
                        model
                    ].get(mlrun.common.schemas.MonitoringData.MODEL_ENDPOINT_UID),
                    mm_schemas.StreamProcessingEvent.LABELS: monitoring_data[model].get(
                        mlrun.common.schemas.MonitoringData.OUTPUTS
                    ),
                    mm_schemas.StreamProcessingEvent.FUNCTION_URI: self.server.function_uri
                    if self.server
                    else None,
                    mm_schemas.StreamProcessingEvent.REQUEST: request,
                    mm_schemas.StreamProcessingEvent.RESPONSE: resp,
                    mm_schemas.StreamProcessingEvent.ERROR: event.body[
                        mm_schemas.StreamProcessingEvent.ERROR
                    ]
                    if mm_schemas.StreamProcessingEvent.ERROR in event.body
                    else None,
                    mm_schemas.StreamProcessingEvent.METRICS: event.body[
                        mm_schemas.StreamProcessingEvent.METRICS
                    ]
                    if mm_schemas.StreamProcessingEvent.METRICS in event.body
                    else None,
                }
            )
        event.body = monitoring_event_list
        return event


class BackgroundTaskStatus(storey.MapClass):
    """
    background task status checker, prevent events from pushing to the model monitoring stream target if model endpoints
    creation failed or in progress
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.server: mlrun.serving.GraphServer = (
            getattr(self.context, "server", None) if self.context else None
        )
        self._background_task_check_timestamp = None
        self._background_task_state = mlrun.common.schemas.BackgroundTaskState.running

    def do(self, event):
        if self.server is None:
            return None
        if (
            self._background_task_state
            == mlrun.common.schemas.BackgroundTaskState.running
            and (
                self._background_task_check_timestamp is None
                or mlrun.utils.now_date() - self._background_task_check_timestamp
                >= timedelta(
                    seconds=mlrun.mlconf.model_endpoint_monitoring.model_endpoint_creation_check_period
                )
            )
        ):
            background_task = mlrun.get_run_db().get_project_background_task(
                self.server.project, self.server.model_endpoint_creation_task_name
            )
            self._background_task_check_timestamp = mlrun.utils.now_date()
            self._log_background_task_state(background_task.status.state)
            self._background_task_state = background_task.status.state

        if (
            self._background_task_state
            == mlrun.common.schemas.BackgroundTaskState.succeeded
        ):
            return event
        else:
            return None

    def _log_background_task_state(
        self, background_task_state: mlrun.common.schemas.BackgroundTaskState
    ):
        logger.info(
            "Checking model endpoint creation task status",
            task_name=self.server.model_endpoint_creation_task_name,
        )
        if (
            background_task_state
            in mlrun.common.schemas.BackgroundTaskState.terminal_states()
        ):
            logger.info(
                f"Model endpoint creation task completed with state {background_task_state}"
            )
        else:  # in progress
            logger.info(
                f"Model endpoint creation task is still in progress with the current state: "
                f"{background_task_state}. Events will not be monitored for the next "
                f"{mlrun.mlconf.model_endpoint_monitoring.model_endpoint_creation_check_period} seconds",
                name=self.name,
                background_task_check_timestamp=self._background_task_check_timestamp.isoformat(),
            )


class SamplingStep(storey.MapClass):
    """sampling step, samples the serving outputs for the model monitoring as sampling_percentage defines"""

    def __init__(
        self,
        sampling_percentage: Optional[float] = 100.0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.sampling_percentage = (
            sampling_percentage if 0 < sampling_percentage <= 100 else 100
        )

    def do(self, event):
        logger.debug(
            "sampling step runs",
            event=event,
            sampling_percentage=self.sampling_percentage,
        )
        if self.sampling_percentage != 100 and not event.get(
            mm_schemas.StreamProcessingEvent.ERROR
        ):
            request = event[mm_schemas.StreamProcessingEvent.REQUEST]
            num_of_inputs = len(request["inputs"])
            sampled_requests_indices = self._pick_random_requests(
                num_of_inputs, self.sampling_percentage
            )
            if not sampled_requests_indices:
                return None

            event[mm_schemas.StreamProcessingEvent.REQUEST]["inputs"] = [
                request["inputs"][i] for i in sampled_requests_indices
            ]

            if isinstance(
                event[mm_schemas.StreamProcessingEvent.RESPONSE]["outputs"], list
            ):
                event[mm_schemas.StreamProcessingEvent.RESPONSE]["outputs"] = [
                    event[mm_schemas.StreamProcessingEvent.RESPONSE]["outputs"][i]
                    for i in sampled_requests_indices
                ]
        event[mm_schemas.EventFieldType.SAMPLING_PERCENTAGE] = self.sampling_percentage
        event[mm_schemas.EventFieldType.EFFECTIVE_SAMPLE_COUNT] = len(
            event.get(mm_schemas.StreamProcessingEvent.REQUEST, {}).get("inputs", [])
        )
        return event

    @staticmethod
    def _pick_random_requests(num_of_reqs: int, percentage: float) -> list[int]:
        """
        Randomly selects indices of requests to sample based on the given percentage

        :param num_of_reqs: Number of requests to select from
        :param percentage: Sample percentage for each request
        :return: A list containing the indices of the selected requests
        """

        return [
            req for req in range(num_of_reqs) if random.random() < (percentage / 100)
        ]


class MockStreamPusher(storey.MapClass):
    def __init__(self, output_stream=None, **kwargs):
        super().__init__(**kwargs)
        stream = self.context.stream if self.context else None
        self.output_stream = output_stream or stream.output_stream

    def do(self, event):
        self.output_stream.push(
            [event], partition_key=mm_schemas.StreamProcessingEvent.ENDPOINT_ID
        )
