from apispec.ext.marshmallow import MarshmallowPlugin
from apispec.ext.marshmallow.openapi import OpenAPIConverter
from apispec.ext.marshmallow.schema_resolver import SchemaResolver


class CustomOpenAPIConverter(OpenAPIConverter):
    """Custom converter."""

    def schema2jsonschema(self, schema):
        if hasattr(schema, "Meta") and hasattr(schema.Meta, "openapi_schema"):
            return schema.Meta.openapi_schema

        super_result = super().schema2jsonschema(schema)
        if hasattr(schema, "Meta") and hasattr(schema.Meta, "example"):
            super_result["example"] = schema.Meta.example
        return super_result

    def resolve_nested_schema(self, schema):
        """Customized method to support OneOfSchema."""
        # using attribute check not to force an extra requirement in projects
        if hasattr(schema, "type_schemas"):
            return {
                "oneOf": [
                    self.resolve_nested_schema(s) for s in schema.type_schemas.values()
                ]
            }

        return super().resolve_nested_schema(schema)


class CustomSchemaResolver(SchemaResolver):
    """Custom resolver."""

    def resolve_schema_dict(self, schema):
        """Customized method to support OneOfSchema."""
        # using attribute check not to force an extra requirement in projects
        if hasattr(schema, "type_schemas"):
            return {
                "oneOf": [
                    self.resolve_schema_dict(s) for s in schema.type_schemas.values()
                ]
            }

        return super().resolve_schema_dict(schema)


class CustomMarshmallowPlugin(MarshmallowPlugin):
    """Plugin with redefined converter and resolver."""

    Converter = CustomOpenAPIConverter
    Resolver = CustomSchemaResolver