utils.py

from __future__ import annotations

import functools
import warnings
from collections.abc import Callable
from typing import Any, TypeVar

from pydantic import BaseModel

from .connection import Connection

__all__ = [
    "ensure_dict",
    "normalize_result",
    "notify_model",
    "request_model",
    "request_model_from_dict",
    "request_optional_model",
    "serialize_params",
    "validate_model",
    "validate_model_from_dict",
    "validate_optional_model",
]

ModelT = TypeVar("ModelT", bound=BaseModel)
MethodT = TypeVar("MethodT", bound=Callable)
ClassT = TypeVar("ClassT", bound=type)
T = TypeVar("T")
MultiParamModelSpec = tuple[type[BaseModel], ...]


def _param_models_name(models: MultiParamModelSpec) -> str:
    return " | ".join(model_type.__name__ for model_type in models)


def _param_models_field_names(models: MultiParamModelSpec) -> tuple[str, ...]:
    shared_fields = set(models[0].model_fields)
    for model_type in models[1:]:
        shared_fields &= set(model_type.model_fields)
    return tuple(field_name for field_name in models[0].model_fields if field_name in shared_fields)


def model_to_kwargs(model_obj: BaseModel, models: MultiParamModelSpec) -> dict[str, Any]:
    kwargs = {
        field_name: getattr(model_obj, field_name)
        for field_name in _param_models_field_names(models)
        if field_name != "field_meta"
    }
    if meta := getattr(model_obj, "field_meta", None):
        kwargs.update(meta)
    return kwargs


def serialize_params(params: BaseModel) -> dict[str, Any]:
    """Return a JSON-serializable representation used for RPC calls."""
    return params.model_dump(by_alias=True, exclude_none=True, exclude_defaults=True)


def normalize_result(payload: Any) -> dict[str, Any]:
    """Convert optional BaseModel/None responses into JSON-friendly payloads."""
    if payload is None:
        return {}
    if isinstance(payload, BaseModel):
        return serialize_params(payload)
    return payload


def ensure_dict(payload: Any) -> dict[str, Any]:
    """Return payload when it is a dict, otherwise an empty dict."""
    return payload if isinstance(payload, dict) else {}


def validate_model(payload: Any, model_type: type[ModelT]) -> ModelT:
    """Validate payload using the provided Pydantic model."""
    return model_type.model_validate(payload)


def validate_model_from_dict(payload: Any, model_type: type[ModelT]) -> ModelT:
    """Validate payload, coercing non-dict values to an empty dict first."""
    return model_type.model_validate(ensure_dict(payload))


def validate_optional_model(payload: Any, model_type: type[ModelT]) -> ModelT | None:
    """Validate payload when it is a dict, otherwise return None."""
    if isinstance(payload, dict):
        return model_type.model_validate(payload)
    return None


async def request_model(
    conn: Connection,
    method: str,
    params: BaseModel,
    response_model: type[ModelT],
) -> ModelT:
    """Send a request with serialized params and validate the response."""
    response = await conn.send_request(method, serialize_params(params))
    return validate_model(response, response_model)


async def request_model_from_dict(
    conn: Connection,
    method: str,
    params: BaseModel,
    response_model: type[ModelT],
) -> ModelT:
    """Send a request and validate the response, coercing non-dict payloads."""
    response = await conn.send_request(method, serialize_params(params))
    return validate_model_from_dict(response, response_model)


async def request_optional_model(
    conn: Connection,
    method: str,
    params: BaseModel,
    response_model: type[ModelT],
) -> ModelT | None:
    """Send a request and validate optional dict responses."""
    response = await conn.send_request(method, serialize_params(params))
    return validate_optional_model(response, response_model)


async def notify_model(conn: Connection, method: str, params: BaseModel) -> None:
    """Send a notification with serialized params."""
    await conn.send_notification(method, serialize_params(params))


def param_model(param_cls: type[BaseModel]) -> Callable[[MethodT], MethodT]:
    """Decorator to map the method parameters to a Pydantic model.
    It is just a marker and does nothing at runtime.
    """

    def decorator(func: MethodT) -> MethodT:
        func.__param_model__ = param_cls  # type: ignore[attr-defined]
        return func

    return decorator


def param_models(*param_cls: type[BaseModel]) -> Callable[[MethodT], MethodT]:
    """Decorator to mark a method as accepting multiple legacy parameter models."""
    if not param_cls:
        raise ValueError("param_models() requires at least one model class")

    def decorator(func: MethodT) -> MethodT:
        func.__param_models__ = param_cls  # type: ignore[attr-defined]
        return func

    return decorator


def to_camel_case(snake_str: str) -> str:
    """Convert snake_case strings to camelCase."""
    components = snake_str.split("_")
    return components[0] + "".join(x.title() for x in components[1:])


def _make_legacy_func(func: Callable[..., T], model: type[BaseModel]) -> Callable[[Any, BaseModel], T]:
    @functools.wraps(func)
    def wrapped(self, params: BaseModel) -> T:
        warnings.warn(
            f"Calling {func.__name__} with {model.__name__} parameter is "  # type: ignore[attr-defined]
            "deprecated, please update to the new API style.",
            DeprecationWarning,
            stacklevel=3,
        )
        kwargs = {
            field_name: getattr(params, field_name) for field_name in model.model_fields if field_name != "field_meta"
        }
        if meta := getattr(params, "field_meta", None):
            kwargs.update(meta)
        return func(self, **kwargs)  # type: ignore[arg-type]

    return wrapped


def _make_compatible_func(func: Callable[..., T], model: type[BaseModel]) -> Callable[..., T]:
    @functools.wraps(func)
    def wrapped(self, *args: Any, **kwargs: Any) -> T:
        param = None
        if not kwargs and len(args) == 1:
            param = args[0]
        elif not args and len(kwargs) == 1:
            param = kwargs.get("params")
        if isinstance(param, model):
            warnings.warn(
                f"Calling {func.__name__} with {model.__name__} parameter "  # type: ignore[attr-defined]
                "is deprecated, please update to the new API style.",
                DeprecationWarning,
                stacklevel=3,
            )
            kwargs = {
                field_name: getattr(param, field_name)
                for field_name in model.model_fields
                if field_name != "field_meta"
            }
            if meta := getattr(param, "field_meta", None):
                kwargs.update(meta)
            return func(self, **kwargs)  # type: ignore[arg-type]
        return func(self, *args, **kwargs)

    return wrapped


def _make_multi_legacy_func(func: Callable[..., T], models: MultiParamModelSpec) -> Callable[[Any, BaseModel], T]:
    model_name = _param_models_name(models)

    @functools.wraps(func)
    def wrapped(self, params: BaseModel) -> T:
        warnings.warn(
            f"Calling {func.__name__} with {model_name} parameter is "  # type: ignore[attr-defined]
            "deprecated, please update to the new API style.",
            DeprecationWarning,
            stacklevel=3,
        )
        return func(self, **model_to_kwargs(params, models))  # type: ignore[arg-type]

    return wrapped


def _make_multi_compatible_func(func: Callable[..., T], models: MultiParamModelSpec) -> Callable[..., T]:
    model_name = _param_models_name(models)

    @functools.wraps(func)
    def wrapped(self, *args: Any, **kwargs: Any) -> T:
        param = None
        if not kwargs and len(args) == 1:
            param = args[0]
        elif not args and len(kwargs) == 1:
            param = kwargs.get("params")
        if isinstance(param, models):
            warnings.warn(
                f"Calling {func.__name__} with {model_name} parameter "  # type: ignore[attr-defined]
                "is deprecated, please update to the new API style.",
                DeprecationWarning,
                stacklevel=3,
            )
            return func(self, **model_to_kwargs(param, models))  # type: ignore[arg-type]
        return func(self, *args, **kwargs)

    return wrapped


def compatible_class(cls: ClassT) -> ClassT:
    """Mark a class as backward compatible with old API style."""
    for attr in dir(cls):
        func = getattr(cls, attr)
        if not callable(func):
            continue
        model = getattr(func, "__param_model__", None)
        models = getattr(func, "__param_models__", None)
        if model is None and models is None:
            continue
        if "_" in attr:
            if models is not None:
                setattr(cls, to_camel_case(attr), _make_multi_legacy_func(func, models))
            else:
                if model is None:
                    continue
                setattr(cls, to_camel_case(attr), _make_legacy_func(func, model))
        else:
            if models is not None:
                setattr(cls, attr, _make_multi_compatible_func(func, models))
            else:
                if model is None:
                    continue
                setattr(cls, attr, _make_compatible_func(func, model))
    return cls