import wagtail.admin.rich_text.editors.draftail.features as draftail_features
from wagtail.admin.rich_text.converters.html_to_contentstate import BlockElementHandler
from wagtail import hooks


@hooks.register("register_rich_text_features")
def register_code_block_feature(features):
    """
    Registering the `code-block` feature, which uses the `CODE-BLOCK` Draft.js block type,
    and is stored as HTML with a <pre><code> tag.
    """
    feature_name = "code-block"
    type_ = "CODE-BLOCK"

    control = {
        "type": type_,
        "label": "</>",
        "description": "Code block",
        "element": "pre",
    }

    features.register_editor_plugin(
        "draftail", feature_name, draftail_features.BlockFeature(control)
    )

    features.register_converter_rule(
        "contentstate",
        feature_name,
        {
            "from_database_format": {
                "pre": BlockElementHandler(type_),
                "pre > code": BlockElementHandler(type_),
            },
            "to_database_format": {
                "block_map": {type_: {"element": "pre", "props": {}}}
            },
        },
    )

    # Optional: add to default features
    features.default_features.append(feature_name)