# Copyright (c) 2024 Airbyte, Inc., all rights reserved.

import re
import sys
from glob import glob
from pathlib import Path

import anyio
import dagger

PYTHON_IMAGE = "python:3.10"
LOCAL_YAML_DIR_PATH = "airbyte_cdk/sources/declarative"
LOCAL_OUTPUT_DIR_PATH = "airbyte_cdk/sources/declarative/models"


PIP_DEPENDENCIES = [
    "datamodel_code_generator==0.26.3",
]


def get_all_yaml_files_without_ext() -> list[str]:
    return [Path(f).stem for f in glob(f"{LOCAL_YAML_DIR_PATH}/*.yaml")]


def generate_init_module_content() -> str:
    header = "# generated by bin/generate_component_manifest_files.py\n"
    for module_name in get_all_yaml_files_without_ext():
        header += f"from .{module_name} import *\n"
    return header


def replace_base_model_for_classes_with_deprecated_fields(post_processed_content: str) -> str:
    """
    Replace the base model for classes with deprecated fields.
    This function looks for classes that inherit from `BaseModel` and have fields marked as deprecated.
    It replaces the base model with `BaseModelWithDeprecations` for those classes.
    """

    # Find classes with deprecated fields
    classes_with_deprecated_fields = set()
    class_matches = re.finditer(r"class (\w+)\(BaseModel\):", post_processed_content)

    for class_match in class_matches:
        class_name = class_match.group(1)
        class_start = class_match.start()
        # Find the next class definition or end of file
        next_class_match = re.search(
            r"class \w+\(",
            post_processed_content[class_start + len(class_match.group(0)) :],
        )
        class_end = (
            len(post_processed_content)
            if next_class_match is None
            else class_start + len(class_match.group(0)) + next_class_match.start()
        )
        class_content = post_processed_content[class_start:class_end]

        # Check if any field has deprecated=True
        if re.search(r"deprecated\s*=\s*True", class_content):
            classes_with_deprecated_fields.add(class_name)

    # update the imports to include the new base model with deprecation warinings
    # only if there are classes with the fields marked as deprecated.
    if len(classes_with_deprecated_fields) > 0:
        # Find where to insert the base model - after imports but before class definitions
        imports_end = post_processed_content.find(
            "\n\n",
            post_processed_content.find("from pydantic.v1 import"),
        )
        if imports_end > 0:
            post_processed_content = (
                post_processed_content[:imports_end]
                + "\n\n"
                + "from airbyte_cdk.sources.declarative.models.base_model_with_deprecations import (\n"
                + "    BaseModelWithDeprecations,\n"
                + ")"
                + post_processed_content[imports_end:]
            )

    # Use the `BaseModelWithDeprecations` base model for the classes with deprecated fields
    for class_name in classes_with_deprecated_fields:
        pattern = rf"class {class_name}\(BaseModel\):"
        replacement = f"class {class_name}(BaseModelWithDeprecations):"
        post_processed_content = re.sub(pattern, replacement, post_processed_content)

    return post_processed_content


async def post_process_codegen(codegen_container: dagger.Container):
    codegen_container = codegen_container.with_exec(
        ["mkdir", "/generated_post_processed"], use_entrypoint=True
    )
    for generated_file in await codegen_container.directory("/generated").entries():
        if generated_file.endswith(".py"):
            original_content = await codegen_container.file(
                f"/generated/{generated_file}"
            ).contents()
            # the space before _parameters is intentional to avoid replacing things like `request_parameters:` with `requestparameters:`
            post_processed_content = original_content.replace(
                " _parameters:", " parameters:"
            ).replace("from pydantic", "from pydantic.v1")

            post_processed_content = replace_base_model_for_classes_with_deprecated_fields(
                post_processed_content
            )

            codegen_container = codegen_container.with_new_file(
                f"/generated_post_processed/{generated_file}", contents=post_processed_content
            )
    return codegen_container


async def main():
    init_module_content = generate_init_module_content()

    async with dagger.Connection(dagger.Config(log_output=sys.stderr)) as dagger_client:
        codegen_container = (
            dagger_client.container()
            .from_(PYTHON_IMAGE)
            .with_exec(["mkdir", "/generated"], use_entrypoint=True)
            .with_exec(["pip", "install", " ".join(PIP_DEPENDENCIES)], use_entrypoint=True)
            .with_mounted_directory(
                "/yaml", dagger_client.host().directory(LOCAL_YAML_DIR_PATH, include=["*.yaml"])
            )
            .with_new_file("/generated/__init__.py", contents=init_module_content)
        )
        for yaml_file in get_all_yaml_files_without_ext():
            codegen_container = codegen_container.with_exec(
                [
                    "datamodel-codegen",
                    "--input",
                    f"/yaml/{yaml_file}.yaml",
                    "--output",
                    f"/generated/{yaml_file}.py",
                    "--disable-timestamp",
                    "--enum-field-as-literal",
                    "one",
                    "--set-default-enum-member",
                    "--use-double-quotes",
                    "--remove-special-field-name-prefix",
                    # allow usage of the extra key such as `deprecated`, etc.
                    "--field-extra-keys",
                    # account the `deprecated` flag provided for the field.
                    "deprecated",
                    # account the `deprecation_message` provided for the field.
                    "deprecation_message",
                ],
                use_entrypoint=True,
            )

        await (
            (await post_process_codegen(codegen_container))
            .directory("/generated_post_processed")
            .export(LOCAL_OUTPUT_DIR_PATH)
        )


anyio.run(main)
