from typing import Any, Generator

import pydantic

from mlflow.exceptions import MlflowException
from mlflow.models.utils import _convert_llm_ndarray_to_list
from mlflow.protos.databricks_pb2 import INTERNAL_ERROR
from mlflow.pyfunc.model import _load_context_model_and_signature
from mlflow.types.responses import (
    ResponsesAgentRequest,
    ResponsesAgentResponse,
    ResponsesAgentStreamEvent,
)
from mlflow.types.type_hints import model_validate


def _load_pyfunc(model_path: str, model_config: dict[str, Any] | None = None):
    context, responses_agent, _ = _load_context_model_and_signature(model_path, model_config)
    return _ResponsesAgentPyfuncWrapper(responses_agent, context)


class _ResponsesAgentPyfuncWrapper:
    """
    Wrapper class that converts dict inputs to pydantic objects accepted by
    :class:`~ResponsesAgent`.
    """

    def __init__(self, responses_agent, context):
        self.responses_agent = responses_agent
        self.context = context

    def get_raw_model(self):
        """
        Returns the underlying model.
        """
        return self.responses_agent

    def _convert_input(self, model_input) -> ResponsesAgentRequest:
        import pandas

        if isinstance(model_input, pandas.DataFrame):
            model_input = {
                k: _convert_llm_ndarray_to_list(v[0])
                for k, v in model_input.to_dict(orient="list").items()
            }
        elif not isinstance(model_input, dict):
            raise MlflowException(
                "Unsupported model input type. Expected a dict or pandas.DataFrame, but got "
                f"{type(model_input)} instead.",
                error_code=INTERNAL_ERROR,
            )
        return ResponsesAgentRequest(**model_input)

    def _response_to_dict(self, response, pydantic_class) -> dict[str, Any]:
        if isinstance(response, pydantic_class):
            return response.model_dump(exclude_none=True)
        try:
            model_validate(pydantic_class, response)
        except pydantic.ValidationError as e:
            raise MlflowException(
                message=(
                    f"Model returned an invalid response. Expected a {pydantic_class.__name__} "
                    f"object or dictionary with the same schema. Pydantic validation error: {e}"
                ),
                error_code=INTERNAL_ERROR,
            ) from e
        return response

    def predict(self, model_input: dict[str, Any], params=None) -> dict[str, Any]:
        """
        Args:
            model_input: A dict with the
                :py:class:`ResponsesRequest <mlflow.types.responses.ResponsesRequest>` schema.
            params: Unused in this function, but required in the signature because
                `load_model_and_predict` in `utils/_capture_modules.py` expects a params field

        Returns:
            A dict with the
            (:py:class:`ResponsesResponse <mlflow.types.responses.ResponsesResponse>`)
            schema.
        """
        request = self._convert_input(model_input)
        response = self.responses_agent.predict(request)
        return self._response_to_dict(response, ResponsesAgentResponse)

    def predict_stream(
        self, model_input: dict[str, Any], params=None
    ) -> Generator[dict[str, Any], None, None]:
        """
        Args:
            model_input: A dict with the
                :py:class:`ResponsesRequest <mlflow.types.responses.ResponsesRequest>` schema.
            params: Unused in this function, but required in the signature because
                `load_model_and_predict` in `utils/_capture_modules.py` expects a params field

        Returns:
            A generator over dicts with the
                (:py:class:`ResponsesStreamEvent <mlflow.types.responses.ResponsesStreamEvent>`)
                schema.
        """
        request = self._convert_input(model_input)
        for response in self.responses_agent.predict_stream(request):
            yield self._response_to_dict(response, ResponsesAgentStreamEvent)