import ast
import importlib
import os
import shutil
from collections.abc import Collection, Generator, Sequence
from copy import deepcopy
from pathlib import Path
from types import ModuleType
from typing import Any

import ast_comments

from cadwyn._package_utils import IdentifierPythonPath, get_package_path_from_module, get_version_dir_path
from cadwyn._utils import get_index_of_latest_schema_dir_in_module_python_path
from cadwyn.codegen._asts import get_all_names_defined_at_toplevel_of_module, read_python_module
from cadwyn.codegen._common import (
    CodegenContext,
    CodegenPlugin,
    GlobalCodegenContext,
    MigrationPlugin,
    PydanticModelWrapper,
    _EnumWrapper,
    _ModuleWrapper,
    get_fields_from_model,
)
from cadwyn.codegen._plugins.class_migrations import class_migration_plugin
from cadwyn.codegen._plugins.class_rebuilding import ClassRebuildingPlugin
from cadwyn.codegen._plugins.class_renaming import ClassRenamingPlugin
from cadwyn.codegen._plugins.import_auto_adding import ImportAutoAddingPlugin
from cadwyn.codegen._plugins.latest_version_aliasing import LatestVersionAliasingPlugin
from cadwyn.codegen._plugins.module_migrations import module_migration_plugin
from cadwyn.structure.versions import Version, VersionBundle

_AUTO_GENERATION_WARNING = "# THIS FILE WAS AUTO-GENERATED BY CADWYN. DO NOT EVER TRY TO EDIT IT BY HAND\n\n"


DEFAULT_CODEGEN_PLUGINS: tuple[CodegenPlugin, ...] = (
    LatestVersionAliasingPlugin(),
    ClassRebuildingPlugin(),
    ClassRenamingPlugin(),
    ImportAutoAddingPlugin(),
)
DEFAULT_CODEGEN_MIGRATION_PLUGINS: tuple[MigrationPlugin, ...] = (
    module_migration_plugin,
    class_migration_plugin,
)


def generate_code_for_versioned_packages(
    template_module: ModuleType,
    versions: VersionBundle,
    *,
    ignore_coverage_for_latest_aliases: bool = True,
    codegen_plugins: Sequence[CodegenPlugin] = DEFAULT_CODEGEN_PLUGINS,
    migration_plugins: Sequence[MigrationPlugin] = DEFAULT_CODEGEN_MIGRATION_PLUGINS,
    extra_context: dict[str, Any] | None = None,
):
    """
    Args:
        template_module: The latest package from which we will generate the versioned packages
        versions: Version bundle to generate versions from
        ignore_coverage_for_latest_aliases: Add a pragma: no cover comment to the star imports in the generated
        version of the latest module.
    """
    extra_context = extra_context or {}

    _generate_versioned_directories(
        template_module,
        versions=list(versions),
        schemas={
            k: PydanticModelWrapper(v, v.__name__, get_fields_from_model(v))
            for k, v in deepcopy(versions.versioned_schemas).items()
        },
        enums={
            k: _EnumWrapper(v, {member.name: member.value for member in v})
            for k, v in deepcopy(versions.versioned_enums).items()
        },
        modules={k: _ModuleWrapper(module) for k, module in versions.versioned_modules.items()},
        extra_context=extra_context | {"ignore_coverage_for_latest_aliases": ignore_coverage_for_latest_aliases},
        codegen_plugins=codegen_plugins,
        migration_plugins=migration_plugins,
    )


def _generate_versioned_directories(
    template_package: ModuleType,
    versions: list[Version],
    schemas: dict[IdentifierPythonPath, PydanticModelWrapper],
    enums: dict[IdentifierPythonPath, _EnumWrapper],
    modules: dict[IdentifierPythonPath, _ModuleWrapper],
    extra_context: dict[str, Any],
    codegen_plugins: Collection[CodegenPlugin],
    migration_plugins: Collection[MigrationPlugin],
):
    for version in versions:
        global_context = GlobalCodegenContext(
            current_version=version,
            versions=versions,
            schemas=schemas,
            enums=enums,
            modules=modules,
            extra=extra_context,
        )
        _generate_directory_for_version(template_package, codegen_plugins, version, global_context)
        for plugin in migration_plugins:
            plugin(global_context)


def _generate_directory_for_version(
    template_package: ModuleType,
    plugins: Collection[CodegenPlugin],
    version: Version,
    global_context: GlobalCodegenContext,
):
    template_dir = get_package_path_from_module(template_package)
    version_dir = get_version_dir_path(template_package, version.value)
    for (
        _relative_path_to_file,
        template_module,
        parallel_file,
    ) in _generate_parallel_directory(
        template_package,
        version_dir,
    ):
        file_source = read_python_module(template_module)
        parsed_file = ast.parse(file_source)
        context = _build_context(global_context, template_dir, version_dir, template_module, parallel_file, parsed_file)

        parsed_file = _apply_module_level_plugins(plugins, parsed_file, context)
        new_module = _apply_per_node_plugins(plugins, parsed_file, context)
        parallel_file.write_text(_AUTO_GENERATION_WARNING + ast_comments.unparse(new_module))


def _apply_module_level_plugins(
    plugins: Collection[CodegenPlugin],
    parsed_file: ast.Module,
    context: CodegenContext,
) -> ast.Module:
    node_type = type(parsed_file)
    for plugin in plugins:
        if issubclass(node_type, plugin.node_type):
            parsed_file = plugin(parsed_file, context)
    return parsed_file


def _apply_per_node_plugins(
    plugins: Collection[CodegenPlugin],
    parsed_file: ast.Module,
    context: CodegenContext,
) -> ast.Module:
    new_body = []

    for node in parsed_file.body:
        node_type = type(node)
        for plugin in plugins:
            if issubclass(node_type, plugin.node_type):
                node = plugin(node, context)  # noqa: PLW2901
        new_body.append(node)

    return ast.Module(body=new_body, type_ignores=[])


def _build_context(
    global_context: GlobalCodegenContext,
    template_dir: Path,
    version_dir: Path,
    template_module: ModuleType,
    parallel_file: Path,
    parsed_file: ast.Module,
):
    if template_module.__name__.endswith(".__init__"):
        module_python_path = template_module.__name__.removesuffix(".__init__")
    else:
        module_python_path = template_module.__name__
    all_names_defined_at_toplevel_of_file = get_all_names_defined_at_toplevel_of_module(
        parsed_file,
        module_python_path,
    )
    index_of_latest_package_dir_in_module_python_path = get_index_of_latest_schema_dir_in_module_python_path(
        template_module,
        template_dir.with_name(version_dir.name),
    )
    return CodegenContext(
        current_version=global_context.current_version,
        versions=global_context.versions,
        schemas=global_context.schemas,
        enums=global_context.enums,
        modules=global_context.modules,
        extra=global_context.extra,
        index_of_latest_package_dir_in_module_python_path=index_of_latest_package_dir_in_module_python_path,
        module_python_path=module_python_path,
        all_names_defined_on_toplevel_of_file=all_names_defined_at_toplevel_of_file,
        template_module=template_module,
        module_path=parallel_file,
    )


def _generate_parallel_directory(
    template_module: ModuleType,
    parallel_dir: Path,
) -> Generator[tuple[Path, ModuleType, Path], Any, None]:
    if template_module.__file__ is None:  # pragma: no cover
        raise ValueError(
            f"You passed a {template_module=} but it doesn't have a file "
            "so it is impossible to generate its counterpart.",
        )
    dir = get_package_path_from_module(template_module)
    parallel_dir.mkdir(exist_ok=True)
    # >>> [cadwyn, structure, schemas]
    template_module_python_path_parts = template_module.__name__.split(".")
    # >>> [home, foo, bar, cadwyn, structure, schemas]
    template_module_path_parts = Path(template_module.__file__).parent.parts
    # >>> [home, foo, bar] = [home, foo, bar, cadwyn, structure, schemas][:-3]
    root_module_path = Path(
        *template_module_path_parts[: -len(template_module_python_path_parts)],
    )
    for subroot, dirnames, filenames in os.walk(dir):
        original_subroot = Path(subroot)
        parallel_subroot = parallel_dir / original_subroot.relative_to(dir)
        if "__pycache__" in dirnames:
            dirnames.remove("__pycache__")
        for dirname in dirnames:
            (parallel_subroot / dirname).mkdir(exist_ok=True)
        for filename in filenames:
            original_file = (original_subroot / filename).absolute()
            parallel_file = (parallel_subroot / filename).absolute()

            if filename.endswith(".py"):
                original_module_path = ".".join(
                    original_file.relative_to(root_module_path).with_suffix("").parts,
                )
                original_module = importlib.import_module(original_module_path)
                yield original_subroot.relative_to(dir), original_module, parallel_file
            else:
                shutil.copyfile(original_file, parallel_file)
