"""
Replace `admin.site.register` with the new `@register` decorator syntax:
https://docs.djangoproject.com/en/stable/releases/1.7/#minor-features
"""

from __future__ import annotations

import ast
from collections.abc import Iterable, MutableMapping
from functools import partial
from typing import Literal, cast
from weakref import WeakKeyDictionary

from tokenize_rt import Offset, Token

from django_upgrade.ast import ast_start_offset
from django_upgrade.data import Fixer, State, TokenFunc
from django_upgrade.tokens import OP, erase_node, extract_indent, insert, reverse_find

fixer = Fixer(
    __name__,
    min_version=(1, 7),
)

# Keep track of classes that could be decorated with `@admin.register()`
# For each class name, store the associated model class names inferred from
# eligible `admin.site.register` calls.


class AdminDetails:
    __slots__ = ("parent", "lineno", "model_names_per_site")

    def __init__(self, parent: ast.AST, lineno: int) -> None:
        self.parent = parent
        self.lineno = lineno
        self.model_names_per_site: dict[str, set[str]] = {}


decorable_admins: MutableMapping[State, dict[str, AdminDetails | None]] = (
    WeakKeyDictionary()
)
# Name of site to set of unregistered model names, or True if potentially all
# models have been unregistered
unregistered_site_models: MutableMapping[State, dict[str, set[str] | Literal[True]]] = (
    WeakKeyDictionary()
)


def _is_django_admin_imported(state: State) -> bool:
    return (
        "admin" in state.from_imports["django.contrib"]
        or "admin" in state.from_imports["django.contrib.gis"]
    )


@fixer.register(ast.ClassDef)
def visit_ClassDef(
    state: State,
    node: ast.ClassDef,
    parents: tuple[ast.AST, ...],
) -> Iterable[tuple[Offset, TokenFunc]]:
    if _is_django_admin_imported(state) and not uses_full_super_in_init_or_new(node):
        admin_detailses = decorable_admins.setdefault(state, {})
        if node.name in admin_detailses:
            # Duplicate name, ignore
            admin_detailses[node.name] = None
            return
        else:
            admin_detailses[node.name] = AdminDetails(parents[-1], node.lineno)

        if not node.decorator_list:
            offset = ast_start_offset(node)
            decorated = False
        else:
            offset = ast_start_offset(node.decorator_list[0])
            decorated = True
        yield (
            offset,
            partial(
                update_class_def,
                name=node.name,
                state=state,
                decorated=decorated,
            ),
        )


def uses_full_super_in_init_or_new(node: ast.ClassDef) -> bool:
    """
    We cannot convert classes using py2 style `super(MyAdmin, self)`
    in the `__init__` or `__new__` method.
    https://docs.djangoproject.com/en/stable/ref/contrib/admin/#the-register-decorator
    """
    visitor = FullSuperVisitor()
    visitor.generic_visit(node)
    return visitor.found_full_super


class FullSuperVisitor(ast.NodeVisitor):
    __slots__ = ("found_full_super",)

    def __init__(self) -> None:
        self.found_full_super = False

    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
        if node.name in ("__init__", "__new__"):
            for subnode in ast.walk(node):
                if (
                    isinstance(subnode, ast.Call)
                    and isinstance(subnode.func, ast.Name)
                    and subnode.func.id == "super"
                    and len(subnode.args) == 2
                ):
                    self.found_full_super = True

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
        # Avoid descending into a new scope
        return None

    def visit_ClassDef(self, node: ast.ClassDef) -> None:
        # Avoid descending into a new scope
        return None


def update_class_def(
    tokens: list[Token], i: int, *, name: str, state: State, decorated: bool
) -> None:
    admin_details = decorable_admins.get(state, {})[name]
    if admin_details is None or not admin_details.model_names_per_site:
        return

    if decorated:
        i = reverse_find(tokens, i, name=OP, src="@")
    j, indent = extract_indent(tokens, i)

    new_src = ""
    for custom_site, model_names in sorted(admin_details.model_names_per_site.items()):
        joined_names = ", ".join(sorted(model_names))
        custom_site_src = f", site={custom_site}" if custom_site else ""
        new_src += f"{indent}@admin.register({joined_names}{custom_site_src})\n"

    insert(tokens, j, new_src=new_src)


@fixer.register(ast.Call)
def visit_Call(
    state: State,
    node: ast.Call,
    parents: tuple[ast.AST, ...],
) -> Iterable[tuple[Offset, TokenFunc]]:
    if (
        _is_django_admin_imported(state)
        and isinstance(parents[-1], ast.Expr)
        and isinstance(node.func, ast.Attribute)
    ):
        if (
            node.func.attr == "register"
            and (
                (  # admin.site.register(...)
                    isinstance(node.func.value, ast.Attribute)
                    and node.func.value.attr == "site"
                    and isinstance(node.func.value.value, ast.Name)
                    and node.func.value.value.id == "admin"
                    and (site_name := "") == ""  # force assign
                )
                or (  # custom_site.register(...)
                    state.looks_like_admin_file
                    and isinstance(node.func.value, ast.Name)
                    and (site_name := node.func.value.id).endswith("site")
                )
            )
            and (
                (
                    len(node.args) == 2
                    and len(node.keywords) == 0
                    and isinstance((admin_arg := node.args[1]), ast.Name)
                )
                or (
                    len(node.args) == 1
                    and len(node.keywords) == 1
                    and node.keywords[0].arg == "admin_class"
                    and isinstance((admin_arg := node.keywords[0].value), ast.Name)
                )
            )
            and (
                isinstance((first_arg := node.args[0]), ast.Name)
                or (
                    isinstance(first_arg, (ast.Tuple, ast.List))
                    and all(isinstance(elt, ast.Name) for elt in first_arg.elts)
                )
            )
        ):
            if isinstance(first_arg, ast.Name):
                model_names = {first_arg.id}
            else:
                # cast() could be removed by using TypeGuard func above
                model_names = {cast(ast.Name, elt).id for elt in first_arg.elts}
            admin_name = admin_arg.id
            admin_details = decorable_admins.get(state, {}).get(admin_name, None)
            unregistered_models = unregistered_site_models.get(state, {}).get(
                site_name, set()
            )
            if (
                unregistered_models is not True
                and not unregistered_models.intersection(model_names)
                and admin_details is not None
                and admin_details.parent == parents[-2]
                and not (site_name and not admin_name.endswith("Admin"))
                and (
                    site_name == ""
                    or (
                        (
                            site_defined_line := get_site_defined_line(
                                parents[0], site_name
                            )
                        )
                        is not None
                        and site_defined_line < admin_details.lineno
                    )
                )
            ):
                admin_details.model_names_per_site.setdefault(site_name, set()).update(
                    model_names
                )
                yield (
                    ast_start_offset(parents[-1]),
                    partial(
                        remove_register,
                        name=admin_name,
                        state=state,
                        node=parents[-1],
                    ),
                )
        elif node.func.attr == "unregister" and (
            (  # admin.site.unregister(...)
                isinstance(node.func.value, ast.Attribute)
                and node.func.value.attr == "site"
                and isinstance(node.func.value.value, ast.Name)
                and node.func.value.value.id == "admin"
                and (site_name := "") == ""  # force assign
            )
            or (  # custom_site.unregister(...)
                state.looks_like_admin_file
                and isinstance(node.func.value, ast.Name)
                and (site_name := node.func.value.id).endswith("site")
            )
        ):
            # potentially all models unregistered, but in some cases we can
            # detect unregistered names
            unregistered_names: set[str] | Literal[True] = True
            if len(node.args) == 1:
                first_arg = node.args[0]
                if isinstance(first_arg, ast.Name):
                    unregistered_names = {first_arg.id}
                elif isinstance(first_arg, (ast.Tuple, ast.List)) and all(
                    isinstance(elt, ast.Name) for elt in first_arg.elts
                ):
                    # argument is a sequence of models
                    unregistered_names = {
                        cast(ast.Name, elt).id for elt in first_arg.elts
                    }

            state_details = unregistered_site_models.get(state, None)
            if state_details is None:
                state_details = {}
                unregistered_site_models[state] = state_details

            if unregistered_names is True:
                state_details[site_name] = True
            else:
                existing_names = state_details.get(site_name)
                if existing_names is None:
                    state_details[site_name] = unregistered_names
                elif existing_names is not True:
                    existing_names.update(unregistered_names)


def remove_register(
    tokens: list[Token], i: int, *, name: str, state: State, node: ast.Expr
) -> None:
    admin_details = decorable_admins.get(state, {})[name]
    if admin_details is None:
        return

    erase_node(tokens, i, node=node)


site_definitions: MutableMapping[ast.Module, dict[str, int | None]] = (
    WeakKeyDictionary()
)


def get_site_defined_line(module: ast.AST, site_name: str) -> int | None:
    assert isinstance(module, ast.Module)
    lines = site_definitions.get(module, None)
    if lines is None:
        lines = {}
        for node in module.body:
            if isinstance(node, ast.ImportFrom):
                for alias in node.names:
                    if alias.asname is not None:
                        name = alias.asname
                    else:
                        name = alias.name

                    if name.endswith("site") and name not in lines:
                        lines[name] = node.lineno
            elif (
                isinstance(node, ast.Assign)
                and len(node.targets) == 1
                and isinstance(node.targets[0], ast.Name)
                and (name := node.targets[0].id) not in lines
            ):
                lines[name] = node.lineno

        site_definitions[module] = lines
    return lines.get(site_name)
