from __future__ import annotations

from typing import Type

from marshmallow import Schema


__all__ = ["response_schema"]


def response_schema(schema: Schema | Type[Schema] | None = None, code: int = 200, description: str | None = None):
    if schema is not None and not isinstance(schema, Schema):
        schema = schema()

    def wrapper(func):
        if not hasattr(func, "__apispec__"):
            setattr(func, "__apispec__", {"request": {}, "responses": {}})

        func.__apispec__["responses"][str(code)] = {
            "schema": schema,
            "description": description or "",
        }

        return func

    return wrapper