# SPDX-License-Identifier: GPL-3.0-or-later
"""
Module for parsing Django models and extracting field information.
"""

import ast
from typing import Dict, List, Tuple


class DjangoModelVisitor(ast.NodeVisitor):
    """
    AST visitor to extract fields from Django models.
    """

    DJANGO_FIELD_TYPES = {
        "Field",
        "CharField",
        "TextField",
        "IntegerField",
        "BooleanField",
        "DateTimeField",
        "ForeignKey",
        "ManyToManyField",
        "OneToOneField",
        "EmailField",
        "URLField",
        "FileField",
        "ImageField",
        "DecimalField",
        "AutoField",
    }

    def __init__(self) -> None:
        self.models: Dict[str, List[str]] = {}
        self.current_model: str | None = None
        self.models_and_blank_fields: Dict[str, List[str]] = {}

    def visit_ClassDef(self, node: ast.ClassDef) -> None:
        """
        Check class definitions, specifically those that inherit from other classes.

        Parameters
        ----------
        node : ast.ClassDef
            A class definition from Python AST (Abstract Syntax Tree).
            It contains information about the class, such as its name, base classes, body, decorators, etc.
        """
        # Only process classes that inherit from something.
        if node.bases:
            self.current_model = node.name
            if self.current_model not in self.models:
                self.models[self.current_model] = []

            self.generic_visit(node)

        self.current_model = None

    def visit_Assign(self, node: ast.Assign) -> None:
        """
        Check assignment statements within a class.

        Parameters
        ----------
        node : ast.Assign
            An assignment definition from Python AST (Abstract Syntax Tree).
            It represents an assignment statement (e.g., x = 42).
        """
        if not self.current_model:
            return

        for target in node.targets:
            if (
                isinstance(target, ast.Name)
                and not target.id.startswith("_")
                and isinstance(node.value, ast.Call)
                and hasattr(node.value.func, "attr")
            ) and any(
                field_type in node.value.func.attr
                for field_type in self.DJANGO_FIELD_TYPES
            ):
                self.models[self.current_model].append(target.id)

                if any(
                    kw.arg == "blank"
                    and isinstance(kw.value, ast.Constant)
                    and kw.value.value is True
                    for kw in node.value.keywords
                ):
                    if self.current_model not in self.models_and_blank_fields:
                        self.models_and_blank_fields[self.current_model] = []

                    self.models_and_blank_fields[self.current_model].append(target.id)


def extract_model_fields(
    models_file: str,
) -> Tuple[Dict[str, List[str]], Dict[str, List[str]]]:
    """
    Extract fields from Django models file.

    Parameters
    ----------
    models_file : str
        A models.py file that defines Django models.

    Returns
    -------
    Tuple(Dict[str, List[str]], Dict[str, List[str]])
        The fields from the models file extracted into a dictionary for future processing.
    """
    with open(models_file, "r", encoding="utf-8") as f:
        content = f.read().strip()
        # Skip any empty lines at the beginning.
        while content.startswith("\n"):
            content = content[1:]

    try:
        tree = ast.parse(content)

    except SyntaxError as e:
        raise SyntaxError(
            f"Failed to parse {models_file}. Make sure it's a valid Python file. Error: {str(e)}"
        ) from e

    visitor = DjangoModelVisitor()
    visitor.visit(tree)

    return visitor.models, visitor.models_and_blank_fields
