import re

from datamodel_code_generator.format import CustomCodeFormatter


class CodeFormatter(CustomCodeFormatter):
    """A class for postprocessing the code generated by datamodel-codegen."""

    def apply(self, code: str) -> str:
        """Applies postprocessing logic to the generated class file.

        Args:
            code: the generated code.

        Returns:
            The code after postprocessing.
        """
        # append 'Properties' to class names in class definitions and type hints
        class_pattern = r"class (\w+)"
        class_names = re.findall(class_pattern, code)

        for class_name in class_names:
            class_name_pattern = rf"\b{class_name}\b"
            code = re.sub(class_name_pattern, f"{class_name}Properties", code)

        # miscellaneous edits
        code = code.replace("DataPackageProperties", "PackageProperties")
        code = code.replace("DataResourceProperties", "ResourceProperties")
        code = code.replace("schema_", "schema")
        code = code.replace("from __future__ import annotations\n", "")
        code = code.replace(r"\N", r"\\N")

        return code
