"""
Contains the base interface that OctoAI endpoints should implement.

Developers that want to create an endpoint should implement the
``Service`` class in this module as directed by the ``octoai`` command-line
interface.
"""
import functools
import inspect
import os
from abc import ABC, abstractmethod
from collections import OrderedDict
from typing import Any, Callable, Mapping, Type

import pydantic_core
from fastapi import Form
from pydantic import BaseModel, Field, create_model

from .types import File

DEFAULT_VOLUME_PATH = "/octoai/cache"


STORE_ASSETS_NOT_OVERRIDDEN = "NOT_OVERRIDDEN"


def volume_path() -> str:
    """Get mounted volume path in docker.

    :return: Docker path.
    """
    docker_path = os.environ.get("OCTOAI_VOLUME_PATH", None)
    if docker_path:
        return docker_path
    else:
        return DEFAULT_VOLUME_PATH


VOLUME_ENVIRONMENT = {
    "HUGGINGFACE_HUB_CACHE": os.path.join(volume_path(), "huggingface_hub_cache"),
    "TORCH_HOME": os.path.join(volume_path(), "torch_home"),
}


class Service(ABC):
    """
    The base interface that OctoAI endpoints should implement.

    Developers that want to create an endpoint should implement this
    class as directed by the ``octoai`` command-line interface.
    """

    def setup(self) -> None:
        """
        Perform service initialization.

        A common operation to include here is loading weights and making
        those available to the ``infer()`` method in a member variable.
        """
        pass

    def store_assets(self) -> None:
        """Download model assets."""
        pass

    @abstractmethod
    def infer(self, **kwargs: Any) -> Any:
        """Perform inference."""
        pass

    setattr(store_assets, STORE_ASSETS_NOT_OVERRIDDEN, True)


class ResponseAnalytics(BaseModel):
    """Additional analytics metadata."""

    inference_time_ms: float = Field(
        description="Inference execution time (without pauses)"
    )
    performance_time_ms: float = Field(
        description="Inference execution time (including pauses)"
    )


def inspect_input_types(service: Service) -> Type[BaseModel]:
    """Create Pydantic input model from ``infer()`` signature."""
    signature = inspect.signature(service.infer)

    if len(signature.parameters) < 1:
        raise ValueError("infer() requires at least one argument")

    return inspect_input_parameters(signature.parameters)


def inspect_input_parameters(
    parameters: Mapping[str, inspect.Parameter]
) -> Type[BaseModel]:
    """Create Pydnatic input mode from the provided signature."""
    args = OrderedDict()

    for arg_name, arg in parameters.items():
        if arg.annotation == inspect.Parameter.empty:
            raise ValueError("infer() requires type annotations for args")

        default = (
            pydantic_core.PydanticUndefined
            if arg.default == inspect.Parameter.empty
            else arg.default
        )
        args[arg_name] = (
            arg.annotation,
            Field(default=default),
        )

    return create_model(
        "Input",
        __config__=None,
        __base__=BaseModel,
        __module__=__name__,
        __validators__=None,
        **args,
    )


def inspect_output_types(method: Callable) -> Type[BaseModel]:
    """Create Pydantic output model from inference method signature."""
    signature = inspect.signature(method)

    if signature.return_annotation == inspect._empty:
        raise ValueError(f"{method.__name__}() requires a return type annotation")

    rets = OrderedDict()
    rets["output"] = (signature.return_annotation, None)
    rets["analytics"] = (ResponseAnalytics, None)

    return create_model(
        "Output",
        __config__=None,
        __base__=BaseModel,
        __module__=__name__,
        __validators__=None,
        **rets,
    )


def implements_form_data(service: Service) -> bool:
    """Check if ``infer_form_data()`` is implemented."""
    for name, method in inspect.getmembers(service, predicate=inspect.ismethod):
        if name == "infer_form_data":
            return True

    return False


def transform_form_data_signature(service: Service) -> inspect.Signature:
    """Add FastAPI Form() defaults to 'str' fields and check compatibility."""
    if not hasattr(service.infer_form_data, "__path__"):  # type: ignore[attr-defined]
        raise ValueError(
            "infer_form_data() must be decorated with @path() "
            "when using form-data inference."
        )

    if service.infer_form_data.__path__ == "/infer":  # type: ignore[attr-defined]
        raise ValueError("Cannot use /infer for form-data inference.")

    params = []
    signature = inspect.signature(service.infer_form_data)  # type: ignore[attr-defined]
    for param in signature.parameters.values():
        if param.annotation == File:
            params.append(param)
        elif param.annotation == str:
            params.append(param.replace(default=Form()))
        else:
            raise ValueError("Only File and str params are supported for form data.")

    if not params:
        raise ValueError("infer_form_data() must have at least one parameter.")

    return inspect.Signature(params)


def path(path: str):
    """Specify the path for a service method."""

    def wrapped(fn: Callable):
        @functools.wraps(fn)
        def wrapped_fn(*args, **kwargs):
            return fn(*args, **kwargs)

        wrapped_fn.__path__ = path  # type: ignore[attr-defined]
        return wrapped_fn

    return wrapped
