ISelectionFilter for both active document and linked document

Revit API
1

Hello everyone,
I have a piece of code implementing an ISelectionFilter that currently allows users to select only Wall and Floor elements from a linked document. Now, I’d like to enhance it so that the filter works flexibly based on the document selected by the user, whether it’s a linked document (as in the current version) or the active document.
In short, I want the ISelectionFilter class to support both scenarios seamlessly.
Can you help me with it?

class ElementSelectionFilter(ISelectionFilter):
    def __init__(self, linkedDoc):
        self.linkedDoc = linkedDoc

    def AllowElement(self, element):
        return True

    def AllowReference(self, reference, point):
        element = self.linkedDoc.GetElement(reference.LinkedElementId)
        if element is not None:
            if element.Category.Id in [
                ElementId(BuiltInCategory.OST_Walls),
                ElementId(BuiltInCategory.OST_Floors)
            ]:               
                return True
        return False