# sourcery skip: snake-case-functions
"""
Merge documentation and type information
- from an doctring-rich and typed stub module
- infor a less well documented and typed stub module
"""
# Copyright Jos Verlinde
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#

import argparse
import re
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Tuple, TypeVar, Union, cast

import libcst as cst
import libcst.matchers as m
from libcst.codemod import CodemodContext, VisitorBasedCodemodCommand
from libcst.codemod.visitors import AddImportsVisitor, GatherImportsVisitor, ImportItem
from libcst.helpers.module import insert_header_comments
from mpflash.logger import log

from stubber.typing_collector import (
    MODULE_KEY,
    AnnoValue,
    StubTypingCollector,
    ensure_mp_available_comment,
    is_deleter,
    is_getter,
    is_property,
    is_setter,
    update_def_docstr,
    update_module_docstr,
)

from .visitors.type_helpers import AddTypeHelpers, GatherTypeHelpers

Mod_Class_T = TypeVar("Mod_Class_T", cst.Module, cst.ClassDef)
"""TypeVar for Module or ClassDef that both support overloads"""

empty_module = cst.parse_module("")  # Debugging aid : empty_module.code_for_node(node)
_code = empty_module.code_for_node


def _is_inspect_placeholder_params(params: str) -> bool:
    """Return whether params contains only MicroPython inspect dummy names."""
    return re.fullmatch(r"(?:(?:self|cls), )?x\d+(?:, x\d+)*", params) is not None


def _is_unknown_params(params: str) -> bool:
    """Return whether params lacks meaningful parameter information."""
    return params in {
        "",
        "...",
        "*args, **kwargs",
        "self",
        "self, *args, **kwargs",
        "cls",
        "cls, *args, **kwargs",
    } or _is_inspect_placeholder_params(params)


def is_decorator(dec: cst.CSTNode, name: str) -> bool:
    """shorthand to determin if something is a specific decorator (simple name only)"""
    return m.matches(dec, m.Decorator(decorator=m.Name(value=name)))


def is_mp_available_decorator(dec: cst.CSTNode) -> bool:
    """True if decorator is mp_available in any form: @mp_available, @mp_available(...), @_pkg.mp_available, @_pkg.mp_available(...)."""
    return m.matches(
        dec,
        m.Decorator(
            decorator=m.OneOf(
                m.Name("mp_available"),
                m.Call(func=m.Name("mp_available")),
                m.Attribute(value=m.DoNotCare(), attr=m.Name("mp_available")),
                m.Call(func=m.Attribute(value=m.DoNotCare(), attr=m.Name("mp_available"))),
            )
        ),
    )


def is_empty_mp_available_call(dec: cst.CSTNode) -> bool:
    """True if decorator is exactly a zero-arg call to mp_available (qualified or not)."""
    if not isinstance(dec, cst.Decorator):
        return False
    inner = dec.decorator
    if not isinstance(inner, cst.Call):
        return False
    if len(inner.args) != 0:
        return False
    func = inner.func
    if isinstance(func, cst.Name):
        return func.value == "mp_available"
    if isinstance(func, cst.Attribute) and isinstance(func.attr, cst.Name):
        return func.attr.value == "mp_available"
    return False


def _typehelper_name(node: cst.CSTNode) -> Optional[str]:
    """Return the simple target name of a type-helper assignment, or None.

    Used to de-duplicate type-helpers (TypeVars, TypeAliases, constants) when
    merging several doc-stubs into a single transform pass.
    """
    if isinstance(node, cst.AnnAssign) and isinstance(node.target, cst.Name):
        return node.target.value
    if isinstance(node, cst.Assign) and len(node.targets) == 1 and isinstance(node.targets[0].target, cst.Name):
        return node.targets[0].target.value
    return None


def _assign_target_name(stmt: cst.CSTNode) -> Optional[str]:
    """Return the target name of a simple assignment statement, or None."""
    if isinstance(stmt, cst.SimpleStatementLine) and len(stmt.body) == 1:
        return _typehelper_name(stmt.body[0])
    return None


def _is_docstring_stmt(stmt: cst.CSTNode) -> bool:
    """Return whether the statement is a bare string expression (docstring)."""
    return (
        isinstance(stmt, cst.SimpleStatementLine)
        and len(stmt.body) == 1
        and isinstance(stmt.body[0], cst.Expr)
        and isinstance(stmt.body[0].value, (cst.SimpleString, cst.ConcatenatedString))
    )


class MergeCommand(VisitorBasedCodemodCommand):
    """
    A libcst transformer that merges the type-rich information from a doc-stub into
    a firmware stub.
    The resulting file will contain information from both sources.

    - module docstring - from source

    - function parameters and types - from docstubs
    - function return types - from docstubs
    - function docstrings - from source

    """

    DESCRIPTION: str = "Merge the type-rich information from a doc-stub into a firmware stub"
    copy_params: bool = True
    copy_docstr: bool = True
    copy_returns: bool = False

    @staticmethod
    def add_args(arg_parser: argparse.ArgumentParser) -> None:
        """Add command-line args that a user can specify for running this codemod."""

        arg_parser.add_argument(
            "--stubfile",
            dest="docstub_file",
            metavar="PATH",
            help="The path to the doc-stub file",
            type=str,
            required=True,
        )

        arg_parser.add_argument(
            "--copy-params",
            dest="copy_params",
            default=False,
        )

        arg_parser.add_argument(
            "--copy-docstr",
            dest="copy_docstr",
            default=True,
        )

        arg_parser.add_argument(
            "--copy-returns",
            dest="copy_returns",
            default=False,
        )

    def __init__(
        self,
        context: CodemodContext,
        docstub_file: Union[Path, str, Sequence[Union[Path, str]]],
        copy_params: bool = False,
        copy_docstr: bool = True,
        copy_returns: bool = False,
    ) -> None:
        """initialize the base class with context, and save our args.

        `docstub_file` may be a single doc-stub path or a sequence of paths. When
        several paths are given their annotations, imports and type-helpers are
        collected and merged, so a single transform pass can enrich the target with
        the information from all of them. This avoids re-parsing a large target file
        once per doc-stub (e.g. machine.pyi enriched from ~23 machine/*.pyi files).
        """
        super().__init__(context)
        self.replace_functiondef_with_classdef = True
        # stack for storing the canonical name of the current function/method
        self.stack: List[str] = []
        # normalize to a list of doc-stub paths
        if isinstance(docstub_file, (str, Path)):
            self.docstub_paths: List[Path] = [Path(docstub_file)]
        else:
            self.docstub_paths = [Path(p) for p in docstub_file]
        # keep the first path for backwards-compatible references / log messages
        self.docstub_path = self.docstub_paths[0] if self.docstub_paths else None

        self.copy_params = copy_params
        self.copy_docstr = copy_docstr
        self.copy_returns = copy_returns

        # store the (merged) annotations
        self.annotations: Dict[
            Tuple[str, ...],  # key: tuple of canonical class/function name
            AnnoValue,
            # Union[TypeInfo, str, List[TypeInfo]],  # value: TypeInfo
        ] = {}
        self.comments: List[str] = []
        self.stub_imports: Dict[str, ImportItem] = {}
        self.all_imports: List[Union[cst.Import, cst.ImportFrom]] = []
        self.type_helpers = {}

        # collect and merge the type information from each doc-stub file
        for docstub_path in self.docstub_paths:
            docstub_source = docstub_path.read_text(encoding="utf-8")
            if not docstub_source:
                continue
            try:
                # parse the doc-stub file
                stub_tree = cst.parse_module(docstub_source)
            except cst.ParserSyntaxError as e:
                # The caller (e.g. enrich_folder) reports this concisely; keep the
                # detailed message at debug level to avoid duplicate error spam when
                # the same malformed doc-stub is used by many boards.
                log.debug(f"Error parsing {docstub_path}: {e}")
                raise ValueError(f"Error parsing {docstub_path}: {e}") from e
            # create fresh collectors for each doc-stub file
            typing_collector = StubTypingCollector()
            import_collector = GatherImportsVisitor(context)
            typevar_collector = GatherTypeHelpers(context)
            # visit the doc-stub file with all collectors
            stub_tree.visit(typing_collector)
            stub_tree.visit(import_collector)
            stub_tree.visit(typevar_collector)
            # merge the collected information (later files win on key conflicts)
            self.annotations.update(typing_collector.annotations)
            self.comments.extend(typing_collector.comments)
            self.stub_imports.update(import_collector.symbol_mapping)
            self.all_imports.extend(import_collector.all_imports)
            # Type-helpers are keyed by scope (e.g. `()` for module level), so many
            # doc-stubs contribute to the same key. Extend (do not overwrite) the
            # per-scope buckets, de-duplicating by target name - this matches the
            # incremental behaviour where each helper was added at most once.
            for scope, helpers in typevar_collector.all_typehelpers.items():
                bucket = self.type_helpers.setdefault(scope, [])
                seen = {n for n in (_typehelper_name(h) for h in bucket) if n}
                for helper in helpers:
                    name = _typehelper_name(helper)
                    if name and name in seen:
                        continue
                    bucket.append(helper)
                    if name:
                        seen.add(name)

    # ------------------------------------------------------------------------

    def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:
        """
        This method is responsible for updating the module node after processing it in the codemod.
        It performs the following tasks:
        1. Adds any needed imports from the doc-stub.
        2. Adds `from module import *` from the doc-stub.
        3. Updates the module docstring.
        4. Updates the comments in the module.

        :param original_node: The original module node.
        :param updated_node: The updated module node after processing.
        :return: The updated module node.
        """
        # --------------------------------------------------------------------
        # add any needed imports from the source doc-stub
        for k in self.stub_imports.keys():
            import_item = self.stub_imports[k]
            if import_item.module_name == self.context.full_module_name:
                # this is an import from the same module we should NOT add it
                continue
            if import_item.module_name.split(".")[0] == self.context.full_module_name and not self.context.filename.endswith(
                "__init__.pyi"
            ):
                # this is an import from a module child module we should NOT add it
                continue
            log.trace(f"add: import {k} = {import_item}")
            AddImportsVisitor.add_needed_import(
                self.context,
                module=import_item.module_name,
                obj=import_item.obj_name,
                asname=import_item.alias,
                relative=import_item.relative,
            )

        # add `from module import *` from the doc-stub
        # FIXME: this cases a problem if there is also a 'from module import foobar' in the firmware stub
        # also all comments get removed from the import
        if self.all_imports:
            for import_item in self.all_imports:
                if isinstance(import_item, cst.ImportFrom):
                    # perhaps this is an import from *
                    if isinstance(import_item.names, cst.ImportStar):
                        # bit of a hack to get the full module name
                        empty_mod = cst.parse_module("")
                        full_module_name = empty_mod.code_for_node(import_item.module)  # type: ignore
                        log.trace(f"add: from {full_module_name} import *")
                        AddImportsVisitor.add_needed_import(
                            self.context,
                            module=full_module_name,
                            obj="*",
                        )
        # --------------------------------------------------------------------
        # Add any typevars to the module
        if self.type_helpers:
            AddTypeHelpers.add_helpers(self.context, self.type_helpers)
            atv = AddTypeHelpers(self.context)
            updated_node = atv.transform_module(updated_node)

        # --------------------------------------------------------------------
        # update the docstring.
        if MODULE_KEY in self.annotations:
            # update/replace  module docstrings
            # todo: or should we add / merge the docstrings?
            docstub_docstr = self.annotations[MODULE_KEY].docstring
            assert isinstance(docstub_docstr, str)
            src_docstr = original_node.get_docstring() or ""
            if src_docstr or docstub_docstr:
                if self.copy_docstr and (docstub_docstr.strip() != src_docstr.strip()):
                    if src_docstr and self.copy_docstr:
                        log.trace(f"Append module docstrings. (new --- old) ")
                        new_docstr = '"""\n' + docstub_docstr + "\n\n---\n" + src_docstr + '\n"""'
                    else:
                        new_docstr = '"""\n' + docstub_docstr + '\n"""'

                    docstr_node = cst.SimpleStatementLine(
                        body=[
                            cst.Expr(
                                value=cst.SimpleString(
                                    value=new_docstr,
                                )
                            )
                        ]
                    )
                    updated_node = update_module_docstr(updated_node, docstr_node)
        # --------------------------------------------------------------------
        # update the comments
        updated_node = insert_header_comments(updated_node, self.comments)

        # --------------------------------------------------------------------
        # make sure that any @overloads that not yet applied  are also added to the firmware stub
        updated_node = self.add_missed_overloads(updated_node, stack_id=())
        # Add any missing @mp_available
        updated_node = self.add_missed_mp_available(updated_node, stack_id=())
        # Add any missing `# mp_available` attributes
        updated_node = self.add_missed_mp_available_attributes(updated_node, stack_id=())
        # Add any missing literal docstrings
        updated_node = self.add_missed_literal_docstrings(updated_node, stack_id=())
        return updated_node

    def add_missed_overloads(self, updated_node: Mod_Class_T, stack_id: tuple) -> Mod_Class_T:
        """
        Add any missing overloads to the updated_node

        """
        missing_overloads = []
        scope_keys = [k for k in self.annotations.keys() if k[: (len(stack_id))] == stack_id]

        for key in scope_keys:
            for overload in self.annotations[key].overloads:
                missing_overloads.append((overload.def_node, key))
            self.annotations[key].overloads = []  # remove for list, assume  works

        if missing_overloads:
            if isinstance(updated_node, cst.Module):
                module_level = True
                updated_body = list(updated_node.body)  # type: ignore
            elif isinstance(updated_node, cst.ClassDef):
                module_level = False
                updated_body = list(updated_node.body.body)  # type: ignore
            else:
                raise ValueError(f"Unsupported node type: {updated_node}")
            # insert each overload just after a function with the same name

            new_classes = []
            for overload, key in missing_overloads:
                matched = False
                matched, i = self.locate_function_by_name(overload, updated_body)
                if matched:
                    log.trace(f"Add @overload for {overload.name.value}")
                    if self.copy_params:
                        docstring_node = self.annotations[key].docstring_node or ""
                        # Use the new overload - but with the existing docstring
                        overload = update_def_docstr(overload, docstring_node)
                    updated_body.insert(i + 1, overload)
                else:
                    # add to the end of the module or  class
                    if module_level and len(key) > 1:
                        # this is a class level overload, for which no class was found
                        class_name = key[0]
                        if class_name not in new_classes:
                            new_classes.append(class_name)
                            # create a class for it, and then add all the overload methods to that class
                            log.trace(f"Add New class  @overload for {overload.name.value} at the end of the module")
                            # create a list of all overloads for this class
                            class_overloads = [overload for overload, k in missing_overloads if k[0] == class_name]
                            class_def = cst.ClassDef(
                                name=cst.Name(value=class_name),
                                body=cst.IndentedBlock(body=class_overloads),
                            )
                            updated_body.append(class_def)
                        else:
                            # already processed this class method
                            pass
                    else:
                        log.trace(f"Add @overload for {overload.name.value} at the end of the class")
                        updated_body.append(overload)

            if isinstance(updated_node, cst.Module):
                updated_node = updated_node.with_changes(body=tuple(updated_body))
            elif isinstance(updated_node, cst.ClassDef):
                b1 = updated_node.body.with_changes(body=tuple(updated_body))
                updated_node = updated_node.with_changes(body=b1)

                # cst.IndentedBlock(body=tuple(updated_body)))  # type: ignore
        return updated_node

    def add_missed_mp_available(self, updated_node: Mod_Class_T, stack_id: tuple) -> Mod_Class_T:
        """
        Add any missing @mp_available to the updated_node

        """
        missing_decorated = []
        scope_keys = [k for k in self.annotations.keys() if k[: (len(stack_id))] == stack_id]

        for key in scope_keys:
            for mpa in self.annotations[key].mp_available:
                # Skip self-referencing: a class's own @mp_available entry should not be
                # added inside itself. It will be consumed at the parent (module) level.
                if key == stack_id and mpa.def_type == "classdef":
                    log.trace(f"Deferring @mp_available classdef to parent scope: key={key} name={mpa.name}")
                    continue
                log.trace(f"Collecting @mp_available: key={key} name={mpa.name} def_type={mpa.def_type}")
                missing_decorated.append((mpa.def_node, key))
            # Only clear if we consumed all items; keep classdef self-refs for parent scope
            remaining = [mpa for mpa in self.annotations[key].mp_available if key == stack_id and mpa.def_type == "classdef"]
            self.annotations[key].mp_available = remaining

        log.trace(
            f"add_missed_mp_available: stack_id={stack_id} missing_decorated={len(missing_decorated)} node_type={type(updated_node).__name__}"
        )

        if missing_decorated:
            if isinstance(updated_node, cst.Module):
                module_level = True
                updated_body = list(updated_node.body)  # type: ignore
            elif isinstance(updated_node, cst.ClassDef):
                module_level = False
                updated_body = list(updated_node.body.body)  # type: ignore
            else:
                raise ValueError(f"Unsupported node type: {updated_node}")
            # insert each decorated function/class just after a function with the same name

            new_classes = []
            for mpa, key in missing_decorated:
                matched = False
                matched, i = self.locate_function_by_name(mpa, updated_body)
                if matched:
                    if isinstance(mpa, cst.ClassDef):
                        # Class already exists in target - skip to avoid duplication
                        log.trace(f"Skip @mp_available class {mpa.name.value} - already exists in target")
                        continue
                    log.trace(f"Add @mp_available for {mpa.name.value}")
                    if self.copy_params and isinstance(mpa, cst.FunctionDef):
                        docstring_node = self.annotations[key].docstring_node or ""
                        # Use the new overload - but with the existing docstring
                        mpa = update_def_docstr(mpa, docstring_node)
                    updated_body.insert(i + 1, mpa)
                else:
                    # add to the end of the module or  class
                    if isinstance(mpa, cst.ClassDef):
                        # This is a full class definition with @mp_available - add it directly
                        log.trace(f"Add @mp_available class {mpa.name.value} at the end of the module/class")
                        updated_body.append(mpa)
                    elif module_level and len(key) > 1:
                        # this is a class level @mp_available, for which no class was found
                        class_name = key[0]
                        if class_name not in new_classes:
                            new_classes.append(class_name)
                            # create a class for it, and then add all the @mp_available methods to that class
                            log.trace(f"Add New class @mp_available for {mpa.name.value} at the end of the module")
                            # create a list of all overloads for this class
                            class_mpas = [mpa for mpa, k in missing_decorated if k[0] == class_name]
                            class_def = cst.ClassDef(
                                name=cst.Name(value=class_name),
                                body=cst.IndentedBlock(body=class_mpas),
                            )
                            updated_body.append(class_def)
                        else:
                            # already processed this class method
                            pass
                    else:
                        log.trace(f"Add @mp_available for {mpa.name.value} at the end of the class")
                        updated_body.append(mpa)

            if isinstance(updated_node, cst.Module):
                updated_node = updated_node.with_changes(body=tuple(updated_body))
            elif isinstance(updated_node, cst.ClassDef):
                b1 = updated_node.body.with_changes(body=tuple(updated_body))
                updated_node = updated_node.with_changes(body=b1)

                # cst.IndentedBlock(body=tuple(updated_body)))  # type: ignore
        return updated_node

    def add_missed_mp_available_attributes(self, updated_node: Mod_Class_T, stack_id: tuple) -> Mod_Class_T:
        """
        Add attribute assignments marked with a `# mp_available` comment that are
        missing from the target. This is the attribute equivalent of the
        `@mp_available` decorator: it force-merges documented module- or class-level
        attributes that createstubs cannot detect on-device.

        When the attribute already exists in the target, the `# mp_available` marker
        is propagated onto it so it survives further source>dest>final merge steps.
        """
        lookup_key = MODULE_KEY if len(stack_id) == 0 else stack_id
        if lookup_key not in self.annotations:
            return updated_node
        mp_attributes = self.annotations[lookup_key].mp_available_attributes
        if not mp_attributes:
            return updated_node

        if isinstance(updated_node, cst.Module):
            updated_body = list(updated_node.body)
        elif isinstance(updated_node, cst.ClassDef):
            updated_body = list(updated_node.body.body)
        else:
            raise ValueError(f"Unsupported node type: {updated_node}")

        # propagate the marker onto attributes already present in the target scope
        existing = set()
        changed = False
        for idx, stmt in enumerate(updated_body):
            name = _assign_target_name(stmt)
            if name and name in mp_attributes:
                existing.add(name)
                marked = ensure_mp_available_comment(stmt)
                if marked is not stmt:
                    log.trace(f"Propagate # mp_available marker to existing attribute {name}")
                    updated_body[idx] = marked
                    changed = True

        # insert after a leading docstring, if any
        insert_at = 1 if updated_body and _is_docstring_stmt(updated_body[0]) else 0

        added = False
        for name, stmt in mp_attributes.items():
            if name in existing:
                continue
            log.trace(f"Add # mp_available attribute {name}")
            updated_body.insert(insert_at, stmt)
            insert_at += 1
            added = True

        if not (added or changed):
            return updated_node

        if isinstance(updated_node, cst.Module):
            updated_node = updated_node.with_changes(body=tuple(updated_body))
        elif isinstance(updated_node, cst.ClassDef):
            b1 = updated_node.body.with_changes(body=tuple(updated_body))
            updated_node = updated_node.with_changes(body=b1)
        return updated_node

    def add_missed_literal_docstrings(self, updated_node: Mod_Class_T, stack_id: tuple) -> Mod_Class_T:
        """
        Add any missing literal docstrings to the updated_node
        """
        # Find literal docstrings for this scope
        literal_docstrings = {}

        # For module level (empty stack_id), use MODULE_KEY
        if len(stack_id) == 0:
            lookup_key = MODULE_KEY
        else:
            lookup_key = stack_id

        if lookup_key in self.annotations and self.annotations[lookup_key].literal_docstrings:
            literal_docstrings = self.annotations[lookup_key].literal_docstrings

        if not literal_docstrings:
            return updated_node

        # Get the body to modify
        if isinstance(updated_node, cst.Module):
            updated_body = list(updated_node.body)
        elif isinstance(updated_node, cst.ClassDef):
            updated_body = list(updated_node.body.body)
        else:
            raise ValueError(f"Unsupported node type: {updated_node}")

        # Find assignment statements and insert docstrings after them
        new_body = []
        i = 0
        while i < len(updated_body):
            stmt = updated_body[i]
            new_body.append(stmt)

            # Check if this is an assignment statement for a literal that has a docstring
            if isinstance(stmt, cst.SimpleStatementLine) and len(stmt.body) == 1:
                literal_name = None

                if isinstance(stmt.body[0], cst.Assign):
                    # Regular assignment: CONST = value
                    targets = stmt.body[0].targets
                    if len(targets) == 1 and isinstance(targets[0].target, cst.Name):
                        literal_name = targets[0].target.value
                elif isinstance(stmt.body[0], cst.AnnAssign):
                    # Annotated assignment: CONST: int = value
                    if isinstance(stmt.body[0].target, cst.Name):
                        literal_name = stmt.body[0].target.value

                # If we found a literal with a docstring, check if docstring is missing
                docstr_node = None
                if literal_name and literal_name in literal_docstrings:
                    # Check if the next statement is already a docstring for this literal
                    has_existing_docstring = False
                    if i + 1 < len(updated_body):
                        next_stmt = updated_body[i + 1]
                        if (
                            isinstance(next_stmt, cst.SimpleStatementLine)
                            and len(next_stmt.body) == 1
                            and isinstance(next_stmt.body[0], cst.Expr)
                            and isinstance(next_stmt.body[0].value, cst.SimpleString)
                        ):
                            has_existing_docstring = True
                            docstr_node = next_stmt
                            i += 1  # skip the docstring node to avoid duplication

                    # If no existing docstring, add the one from doc_stub
                    if not has_existing_docstring or self.copy_docstr:
                        docstr_node = literal_docstrings[literal_name]
                        log.trace(f"Added literal docstring for {literal_name}")
                    if docstr_node:
                        new_body.append(docstr_node)

            i += 1

        # Update the node with the new body
        if isinstance(updated_node, cst.Module):
            updated_node = updated_node.with_changes(body=tuple(new_body))
        elif isinstance(updated_node, cst.ClassDef):
            new_indented_body = updated_node.body.with_changes(body=tuple(new_body))
            updated_node = updated_node.with_changes(body=new_indented_body)

        return updated_node

    def locate_function_by_name(self, overload, updated_body):
        """locate the (last) function or class by name"""
        matched = False
        i = 0
        for i, node in reversed(list(enumerate(updated_body))):
            if isinstance(node, (cst.FunctionDef, cst.ClassDef)) and node.name.value == overload.name.value:
                matched = True
                break
        return matched, i
        # --------------------------------------------------------------------

    # ------------------------------------------------------------

    def visit_ClassDef(self, node: cst.ClassDef) -> Optional[bool]:
        """keep track of the the (class, method) names to the stack"""
        self.stack.append(node.name.value)

    def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.ClassDef:
        stack_id = tuple(self.stack)
        self.stack.pop()
        if stack_id not in self.annotations:
            # no changes to the class
            return updated_node
        # update the firmware_stub from the doc_stub information
        doc_stub = self.annotations[stack_id].type_info
        # first update the docstring
        updated_node = update_def_docstr(updated_node, doc_stub.docstr_node)
        # Sometimes the firmware stubs and the doc stubs have different types : FunctionDef / ClassDef
        # we need to be careful not to copy over all the annotations if the types are different
        if doc_stub.def_type == "classdef":
            # Same type, we can copy over all the annotations
            # combine the decorators from the doc-stub and the firmware stub, deduplicating @mp_available
            new_decorators = []
            if doc_stub.decorators:
                new_decorators.extend(doc_stub.decorators)
            if updated_node.decorators:
                for dec in updated_node.decorators:
                    # Skip @mp_available decorators already present from the doc-stub
                    if is_mp_available_decorator(dec) and any(is_mp_available_decorator(d) for d in new_decorators):
                        continue
                    new_decorators.append(dec)

            updated_node = updated_node.with_changes(
                decorators=new_decorators,
                bases=doc_stub.def_node.bases,  # type: ignore
            )
        # else:
        # Different type: ClassDef != FuncDef ,
        # for now just return the updated node
        # Add any missing methods overloads
        updated_node = self.add_missed_overloads(updated_node, stack_id)
        # Add any missing @mp_available
        updated_node = self.add_missed_mp_available(updated_node, stack_id)
        # Add any missing `# mp_available` attributes
        updated_node = self.add_missed_mp_available_attributes(updated_node, stack_id)
        # Add any missing literal docstrings
        updated_node = self.add_missed_literal_docstrings(updated_node, stack_id)
        return updated_node

    # ------------------------------------------------------------------------

    # ------------------------------------------------------------------------
    def visit_FunctionDef(self, node: cst.FunctionDef) -> Optional[bool]:
        self.stack.append(node.name.value)
        return True

    def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> Union[cst.FunctionDef, cst.ClassDef]:
        "Update the function Parameters and return type, decorators and docstring"
        if is_getter(updated_node):
            extra = ["getter"]
        elif is_setter(updated_node):
            extra = ["setter"]
        elif is_deleter(updated_node):
            extra = ["deleter"]
        else:
            extra = []
        stack_id = tuple(self.stack + extra)
        self.stack.pop()
        if stack_id not in self.annotations:
            # no changes to the function in docstub
            return updated_node
        if updated_node.decorators and any(is_decorator(dec, "overload") for dec in updated_node.decorators):
            # do not overwrite existing @overload functions
            # Consume its source counterpart so add_missed_overloads does not
            # append the complete overload group again on repeated merges.
            overloads = self.annotations[stack_id].overloads
            for index, overload in enumerate(overloads):
                params_match = overload.params is not None and updated_node.params.deep_equals(overload.params)
                returns_match = (updated_node.returns is None and overload.returns is None) or (
                    updated_node.returns is not None and overload.returns is not None and updated_node.returns.deep_equals(overload.returns)
                )
                if params_match and returns_match:
                    overloads.pop(index)
                    break
            return updated_node
        # Force-merge when firmware stub marks with @mp_available() (no args)
        if updated_node.decorators and any(is_empty_mp_available_call(dec) for dec in updated_node.decorators):
            if self.annotations[stack_id].mp_available:
                doc_stub = self.annotations[stack_id].mp_available.pop(0)
                # Clear remaining to avoid later duplication via add_missed_mp_available
                self.annotations[stack_id].mp_available = []
            else:
                doc_stub = self.annotations[stack_id].type_info
            log.debug(f"Force merge via @mp_available() :{updated_node.name.value}")
            updated_node = self.merge_decorator(original_node, updated_node, stack_id, doc_stub)
            return updated_node

        # update the firmware_stub from the doc_stub information
        doc_stub = self.annotations[stack_id].type_info
        if isinstance(doc_stub.def_node, cst.FunctionDef):
            # avoid mismatching property decorators
            # if the updated node is a property, and the doc_stub node is not a property
            # then we should not update the node
            if is_property(updated_node) and not is_property(doc_stub.def_node):
                return updated_node
            if is_setter(updated_node) and not is_setter(doc_stub.def_node):
                return updated_node
            if is_getter(updated_node) and not is_getter(doc_stub.def_node):
                return updated_node
            if is_deleter(updated_node) and not is_deleter(doc_stub.def_node):
                return updated_node

        # Check if it is an @overload decorator
        add_overload = any(is_decorator(dec, "overload") for dec in doc_stub.decorators) and len(self.annotations[stack_id].overloads) >= 1

        # If there are overloads in the documentation, use the first one
        if add_overload:
            log.debug(f"Change to @overload :{updated_node.name.value}")
            # Use the new overload - but with the existing docstring
            doc_stub = self.annotations[stack_id].overloads.pop(0)
            updated_node = self.merge_decorator(original_node, updated_node, stack_id, doc_stub)
            return updated_node

        # If there are mp_available decorators in the documentation, use the first one
        add_mpa_deco = (
            any(is_mp_available_decorator(dec) for dec in doc_stub.decorators) and len(self.annotations[stack_id].mp_available) >= 1
        )
        if add_mpa_deco:
            log.debug(f"Change to @mp_available :{updated_node.name.value}")
            # Use the new @mp_available - but with the existing docstring
            doc_stub = self.annotations[stack_id].mp_available.pop(0)
            updated_node = self.merge_decorator(original_node, updated_node, stack_id, doc_stub)
            return updated_node

        # assert isinstance(doc_stub, TypeInfo)
        # assert doc_stub
        # first update the docstring
        has_no_docstring = updated_node.get_docstring() is None
        if has_no_docstring or self.copy_docstr:
            # overwrite existing docstring if there was none , or if it is asked
            updated_node = update_def_docstr(updated_node, doc_stub.docstr_node, doc_stub.def_node)

        # Sometimes the firmware stubs and the doc stubs have different types : FunctionDef / ClassDef
        # we need to be careful not to copy over all the annotations if the types are different
        if doc_stub.def_type == "funcdef":
            # Same type, we can copy over the annotations
            # params that should  not be overwritten by the doc-stub ?
            if self.copy_params:
                # we are copying rich type definitions, just assume they are better than what is currently
                # in the destination stub
                overwrite_params = True
            else:
                params_txt = _code(original_node.params)
                overwrite_params = _is_unknown_params(params_txt)

            # return that should not be overwritten by the doc-stub ?
            # copy_returns=True explicitly prefers source return annotations.
            overwrite_return = self.copy_returns
            if not overwrite_return:
                overwrite_return = True
                if original_node.returns:
                    try:
                        overwrite_return = original_node.returns.annotation.value in [  # type: ignore
                            "Incomplete",
                            "Any",
                            "...",
                        ]
                    except AttributeError:
                        pass
            # combine the decorators from the doc-stub and the firmware stub
            new_decorators = []
            if doc_stub.decorators:
                new_decorators.extend(doc_stub.decorators)

            for decorator in updated_node.decorators:
                if _code(decorator) not in [_code(d) for d in new_decorators]:
                    new_decorators.append(decorator)

            if len(stack_id) == 1:
                new_decorators = [
                    decorator
                    for decorator in new_decorators
                    if not (is_decorator(decorator, "classmethod") or is_decorator(decorator, "staticmethod"))
                ]

            # if the method is both a static and a class method, we remove the @classmethod decorator to avoid inconsistencies
            if any(is_decorator(dec, "staticmethod") for dec in new_decorators) and any(
                is_decorator(dec, "classmethod") for dec in new_decorators
            ):
                new_decorators = [dec for dec in new_decorators if dec.decorator.value != "classmethod"]

            source_asynchronous = doc_stub.def_node.asynchronous if isinstance(doc_stub.def_node, cst.FunctionDef) else None
            return updated_node.with_changes(
                asynchronous=source_asynchronous or updated_node.asynchronous,
                decorators=new_decorators,
                params=doc_stub.params if overwrite_params else updated_node.params,
                returns=doc_stub.returns if overwrite_return else updated_node.returns,
            )

        elif doc_stub.def_type == "classdef":
            # Different type: ClassDef != FuncDef ,
            if doc_stub.def_node and self.replace_functiondef_with_classdef:
                # replace the functionDef with the classdef from the stub file
                return doc_stub.def_node
            # for now just return the updated node
            return updated_node
        else:
            #  just return the updated node
            return updated_node

    def merge_decorator(
        self,
        original_node: cst.FunctionDef,
        updated_node: cst.FunctionDef,
        stack_id,
        doc_stub,
    ):
        if not self.copy_params:
            # we have copied over the entire function definition, no further processing should be done on this node
            doc_stub.def_node = cast(cst.FunctionDef, doc_stub.def_node)
            updated_node = doc_stub.def_node

        else:
            # Save (first) existing docstring if any
            existing_ds = None
            if updated_node.get_docstring():
                # if there is one , then get it including the layout
                existing_ds = original_node.body.body[0]
                assert isinstance(existing_ds, cst.SimpleStatementLine)

            self.annotations[stack_id].docstring_node = existing_ds
            updated_node = update_def_docstr(doc_stub.def_node, existing_ds)
        return updated_node
