import importlib
import inspect
import logging
import os
import pkgutil
import sys
import warnings
from typing import List, Optional

from redbaron import RedBaron

from fixers import fix_access_controls, fix_imports, fix_xcom_pull, fix_xcom_push

# logging.basicConfig(level=logging.DEBUG)

warnings.filterwarnings("error")

LOG_FORMAT = "[%(name)s] [%(levelname)s] %(message)s"
LOG_LEVEL = logging.INFO
formatter = logging.Formatter(LOG_FORMAT)
nova_logger = logging.getLogger("nova")
nova_logger.setLevel(LOG_LEVEL)
handler = logging.StreamHandler()
handler.setLevel(LOG_LEVEL)
handler.setFormatter(formatter)
nova_logger.addHandler(handler)

# TODO Give control over fixers to run
_fixers = [
    "fix_imports",
    "fix_access_controls",
    "fix_xcom_push",
    "fix_xcom_pull",
    "update_dag_name",
    # TODO Worker Queues
]
_reports = [
    # TODO Async Optimizations
    # TODO Worker Queues
]

common_dag_class_imports = [
    "airflow.hooks",
    "airflow.operators",
    "airflow.sensors",
    "airflow.providers",
]


def _get_class_import_map(root_module) -> dict:
    found_packages = list(pkgutil.iter_modules(root_module.__path__))
    if not found_packages:
        return {}
    # action_names = [
    #     (name.split(".")[-1], ispkg) for _, name, ispkg in found_packages
    # ]
    # importer = found_packages[0][0]
    class_import_map = {}
    for x, name, ispkg in found_packages:
        if ispkg:
            try:
                module = importlib.import_module(f"{root_module.__name__}.{name}")
            except Exception as e:
                breakpoint()
            submodule_class_import_map = _get_class_import_map(module)
            class_import_map = {**submodule_class_import_map, **class_import_map}

        else:
            try:
                module = importlib.import_module(f"{root_module.__name__}.{name}")
                if f"{root_module.__name__}.{name}" not in sys.modules:
                    continue
                classes = inspect.getmembers(module, inspect.isclass)
                for class_name, c in classes:
                    if c.__module__ != f"{root_module.__name__}.{name}":
                        continue
                    class_import_map[class_name] = {
                        "value": f"{root_module.__name__}.{name}",
                        "targets": class_name,
                    }
            except DeprecationWarning as err:
                nova_logger.debug(
                    f"Skip deprecated module {root_module.__name__}.{name}: {err}"
                )
            except ModuleNotFoundError as err:
                nova_logger.error(
                    f"Cannot load modules from {root_module.__name__}.{name}: {err}"
                )
            except UserWarning as err:
                continue
    return class_import_map


def _generate_import_map() -> dict:
    class_import_map = {}
    for common_dag_class_import in common_dag_class_imports:
        root_module = importlib.import_module(common_dag_class_import)
        class_import_map = {
            **_get_class_import_map(root_module),
            **class_import_map,
        }
    return class_import_map


def _upgrade_dag_file(
    filename: str, out_dir: str, class_import_map, rename_dag: Optional[bool] = False
) -> None:
    if "upgraded" in filename:
        return
    red = _load_dag_file(filename)
    fix_imports(class_import_map, red)
    fix_access_controls(red)
    fix_xcom_push(red)
    fix_xcom_pull(red)
    if rename_dag:
        _update_dag_name(red)
    _write_updated_dag_file(filename, out_dir, red)


def _load_dag_file(filename: str) -> RedBaron:
    with open(filename, "r") as f:
        code = f.read()
    return RedBaron(code)


def _write_updated_dag_file(filename: str, out_dir: str, red: RedBaron) -> None:
    upgraded_filename = filename
    upgraded_filename = os.path.join(out_dir, os.path.basename(upgraded_filename))

    with open(upgraded_filename, "w") as f:
        f.write(red.dumps())


def _update_dag_name(red: RedBaron) -> None:
    # update dag name
    dag_def = red.find("NameNode", value="DAG").parent
    dag_name = dag_def.call_[0].string_.value
    file_name = dag_name[:-1]
    ext = dag_name[-1]
    dag_def.call_[0].value = f"{file_name}_upgraded{ext}"


def _create_out_dir_if_not_exists(out_dir: str) -> None:
    os.makedirs(out_dir, exist_ok=True)


def _get_dag_files_from_dir(dag_dir: str) -> List[str]:
    return [
        os.path.join(dag_dir, f)
        for f in os.listdir(dag_dir)
        if os.path.isfile(os.path.join(dag_dir, f)) and f[-3:] == ".py"
    ]


def upgrade_dag_files(
    dag_dir: str,
    out_dir: Optional[str] = None,
    rename_dag: Optional[bool] = False,
):
    if not dag_dir:
        nova_logger.error("No DAG directory provided.")
        exit()
    if not out_dir:
        out_dir = os.path.join(dag_dir, "upgraded/")

    _create_out_dir_if_not_exists(out_dir)
    class_import_map = _generate_import_map()

    dag_files = _get_dag_files_from_dir(dag_dir)
    for dag_file in dag_files:
        _upgrade_dag_file(dag_file, out_dir, class_import_map, rename_dag)


def upgrade_dag_file(
    dag_file: str,
    out_dir: Optional[str] = None,
    rename_dag: Optional[bool] = False,
):
    _create_out_dir_if_not_exists(out_dir)
    class_import_map = _generate_import_map()
    _upgrade_dag_file(dag_file, out_dir, class_import_map, rename_dag)
