#!/usr/bin/env python3

"""Find view definition files and execute them."""

from argparse import ArgumentParser
import os
from pathlib import Path
import re
import string
import yaml

UDF_DIRS = ("udf/", "udf_js/")
UDF_DATASETS = ["udf", "udf_js"]
MOZFUN_DIR = "mozfun"
DESCRIPTION_RE = re.compile(r"(/\*([^*]|[\r\n]|(\*+([^*/]|[\r\n])))*\*+/)|(//.*)")
UDF_RE = r"((:?udf|udf_js)\.[a-zA-z0-9_]+)"
UDF_DEFINITION = r"((:?udf|udf_js)\.[a-zA-z0-9_]+\([^\)]*\))"

parser = ArgumentParser(description=__doc__)
parser.add_argument(
    "--udf", help="Migrate the specified UDF to mozfun.",
)
parser.add_argument(
    "--udf-dirs",
    "--udf_dirs",
    nargs="+",
    default=UDF_DIRS,
    help="Directories containing UDFs to migrate",
)


def migrate_udf(udf_file):
    """Migrates a single UDF to mozfun."""
    # get the target dataset
    dataset = input("Target dataset name: ")

    file_name = os.path.basename(udf_file)

    old_udf_name = file_name.replace(".sql", "")
    udf_name = old_udf_name.replace(dataset + "_", "")

    custom_udf_name = input(f"Custom UDF name (default: {udf_name}): ")

    friendly_name = string.capwords(udf_name.replace("_", " "))

    if custom_udf_name != "":
        udf_name = custom_udf_name

    description = ""
    old_udf_definition = ""

    with open(udf_file) as udf:
        udf_content = udf.read()
        old_udf_definition = re.findall(UDF_DEFINITION, udf_content.replace("\n", ""))[
            0
        ][0]

        for udf_dataset in UDF_DATASETS:
            # replace shared-prod UDF references
            # we'll assume that referenced UDFs are also getting exported to mozfun
            udf_content = udf_content.replace(
                f"{udf_dataset}.{old_udf_name}", f"{dataset}.{udf_name}"
            )

            udf_references = re.findall(UDF_RE, udf_content)
            for udf_reference in set(udf_references):
                custom_udf_reference = input(
                    f"Referenced {udf_reference[0]} should be: "
                )

                if custom_udf_reference != "":
                    udf_content = udf_content.replace(
                        udf_reference[0], custom_udf_reference
                    )

        comment = re.findall(DESCRIPTION_RE, udf_content)

        if len(comment) > 0:
            description = comment[0][0].replace("/*", "")
            description = description.replace("*/", "").strip()
            description = description.replace("\n", " ").strip()

        # write files to mozfun/
        mozfun_udf_dir = Path(MOZFUN_DIR) / dataset / udf_name
        mozfun_udf_dir.mkdir(parents=True, exist_ok=True)

        Path(mozfun_udf_dir / "udf.sql").write_text(udf_content)

        metadata_content = {
            "description": description,
            "friendly_name": friendly_name,
        }
        Path(mozfun_udf_dir / "metadata.yaml").write_text(yaml.dump(metadata_content))

    # update shared-prod UDF to reference mozfun UDF
    with open(udf_file, "w") as wrapper_udf:
        mozfun_udf = old_udf_definition

        for udf_dataset in UDF_DATASETS:
            mozfun_udf = mozfun_udf.replace(
                f"{udf_dataset}.{old_udf_name}", f"{dataset}.{udf_name}"
            )
            mozfun_udf = re.sub(r"[^\s-]+,", ",", mozfun_udf)
            mozfun_udf = re.sub(r"[^\s-]+\)", ")", mozfun_udf)

        wrapper_content = (
            "-- Legacy wrapper around a function moved to mozfun.\n"
            + f"CREATE OR REPLACE FUNCTION {old_udf_definition} AS (\n"
            + f"  mozfun.{mozfun_udf}\n"
            + ");\n"
        )

        wrapper_udf.write(wrapper_content)
        wrapper_udf.close()


def main():
    args = parser.parse_args()

    if args.udf:
        # migrate a single UDF
        migrate_udf(args.udf)
        return

    # iterate through udfs and migrate one by one
    for udf_dir in args.udf_dirs:
        if os.path.isdir(udf_dir):
            for root, dirs, files in os.walk(udf_dir):
                for udf_file in files:
                    if udf_file.endswith(".sql"):
                        with open(os.path.join(root, udf_file)) as udf:
                            if (
                                "Legacy wrapper around a function moved to mozfun"
                                not in udf.read()
                            ):
                                resp = input(f"Migrate {udf_file}? [y/n] ")

                                if resp == "y" or resp == "yes":
                                    migrate_udf(os.path.join(root, udf_file))


if __name__ == "__main__":
    main()
