from typing import TYPE_CHECKING

if TYPE_CHECKING:
    pass


def _get_obj_id(obj) -> int:
    """Extracts a non-zero indirect PDF object ID, returning 0 for direct or invalid objects."""
    try:
        objgen = getattr(obj, "objgen", None)
        if objgen and objgen[0] != 0:
            return int(objgen[0])
    except (AttributeError, TypeError, IndexError):
        pass
    return 0


def get_xobject_ocg_ids(xobj) -> set:
    """Extracts a set of OCG object IDs from an XObject's /OC dictionary."""
    import pikepdf

    if "/OC" not in xobj:
        return set()

    oc = xobj.OC
    if not isinstance(oc, pikepdf.Dictionary):
        return set()

    # An OCMD can contain a single OCG or an array of OCGs
    if oc.get("/Type") == "/OCMD":
        return _ocg_ids_from_ocmd(oc.get("/OCGs"), pikepdf.Array)

    obj_id = _get_obj_id(oc)
    if obj_id:
        return {obj_id}

    return set()


def _ocg_ids_from_ocmd(ocgs, pikepdf_array):
    ocg_ids = set()
    if isinstance(ocgs, pikepdf_array):
        for o in ocgs:
            obj_id = _get_obj_id(o)
            if obj_id:
                ocg_ids.add(obj_id)
    else:
        obj_id = _get_obj_id(ocgs)
        if obj_id:
            ocg_ids.add(obj_id)
    return ocg_ids


def get_page_layer_map(resources) -> tuple[dict, dict]:
    """
    Analyzes a resource dictionary and returns maps for both inline properties and XObjects.
    Returns: (properties_map, xobject_map)
    """
    import pikepdf

    prop_map: dict[str, int] = {}
    xobj_map: dict[str, set] = {}

    if resources is None:
        return prop_map, xobj_map

    # 1. Map /Properties (Used for inline BDC tags)
    if "/Properties" in resources:
        _copy_to_prop_map(resources.Properties, prop_map, pikepdf.Dictionary)

    # 2. Map /XObject (Used for 'Do' commands)
    if "/XObject" in resources:
        _copy_to_xobj_map(resources.XObject, xobj_map)

    return prop_map, xobj_map


def _copy_to_prop_map(resources_properties, prop_map, pikepdf_dictionary):
    for local_name, pdf_obj in resources_properties.items():
        if isinstance(pdf_obj, pikepdf_dictionary) and pdf_obj.get("/Type") == "/OCG":
            obj_id = _get_obj_id(pdf_obj)
            if obj_id:
                prop_map[str(local_name)] = obj_id


def _copy_to_xobj_map(resources_xobject, xobj_map):
    for local_name, xobj in resources_xobject.items():
        ocg_ids = get_xobject_ocg_ids(xobj)
        if ocg_ids:
            xobj_map[str(local_name)] = ocg_ids


def _remove_targets_from_array(node, target_ids: set):
    """Recursively removes objects matching the target IDs from a pikepdf Array."""
    import pikepdf

    for i in range(len(node) - 1, -1, -1):
        item = node[i]
        if isinstance(item, pikepdf.Array):
            _remove_targets_from_array(item, target_ids)
            if len(item) == 0:
                del node[i]
        else:
            obj_id = _get_obj_id(item)
            if obj_id and obj_id in target_ids:
                del node[i]


def clean_ocproperties(pdf, target_ids: set):
    """Safely purges stripped/flattened layers from the PDF's global metadata."""
    from pikepdf import NamePath

    if "/OCProperties" not in pdf.Root:
        return

    ocgs = _clean_master_ocg_array(pdf, target_ids, NamePath)
    _clean_default_config(pdf, target_ids, NamePath)
    _clean_alternate_configs(pdf, target_ids, NamePath)
    _clean_empty_shell(pdf, ocgs)


def _clean_master_ocg_array(pdf, target_ids, pikepdf_namepath):
    # 1. Clean Master OCGs Array
    ocgs = pdf.Root.get(pikepdf_namepath.OCProperties.OCGs)
    if ocgs is not None:
        _remove_targets_from_array(ocgs, target_ids)
    return ocgs


def _clean_default_config(pdf, target_ids, pikepdf_namepath):
    # 2. Clean Default Config
    for key in ["/ON", "/OFF", "/Order"]:
        arr = pdf.Root.get(pikepdf_namepath("/OCProperties", "/D", key))
        if arr is not None:
            _remove_targets_from_array(arr, target_ids)


def _clean_alternate_configs(pdf, target_ids, pikepdf_namepath):
    # 3. Clean Alternate Configs
    configs = pdf.Root.get(pikepdf_namepath.OCProperties.Configs)
    if configs is not None:
        for config in configs:
            for key in ["/ON", "/OFF", "/Order"]:
                arr = config.get(key)
                if arr is not None:
                    _remove_targets_from_array(arr, target_ids)


def _clean_empty_shell(pdf, ocgs):
    # 4. If NO layers are left, completely destroy the OCProperties shell
    if ocgs is not None and len(ocgs) == 0:
        del pdf.Root["/OCProperties"]


def create_layer(pdf, layer_name: str):
    """
    Creates a new Optional Content Group (layer) and registers it globally
    in the PDF's /OCProperties. Returns the OCG object.
    """
    from pikepdf import Array, Dictionary, Name

    # 1. Create the base Optional Content Group (OCG)
    ocg = pdf.make_indirect(Dictionary(Type=Name.OCG, Name=layer_name))

    # 2. Ensure the global /OCProperties dictionary exists
    if "/OCProperties" not in pdf.Root:
        pdf.Root.OCProperties = Dictionary(OCGs=Array(), D=Dictionary(Order=Array(), ON=Array()))

    oc_props = pdf.Root.OCProperties

    # 3. Safely append to the master OCG list
    if "/OCGs" not in oc_props:
        oc_props.OCGs = Array()
    oc_props.OCGs.append(ocg)

    # 4. Safely append to the Default view dictionary (/D) and Order array
    if "/D" not in oc_props:
        oc_props.D = Dictionary(Order=Array(), ON=Array())
    if "/Order" not in oc_props.D:
        oc_props.D.Order = Array()

    oc_props.D.Order.append(ocg)

    # Best practice: ensure it is toggled ON by default
    if "/ON" not in oc_props.D:
        oc_props.D.ON = Array()
    oc_props.D.ON.append(ocg)

    return ocg


def set_layer_state(pdf, target_ids: set, action: str):
    """
    Applies state changes (show/hide/lock/unlock) to the global /OCProperties.
    """
    import pikepdf

    if "/OCProperties" not in pdf.Root or "/D" not in pdf.Root.OCProperties:
        return

    d_dict = pdf.Root.OCProperties.D

    def _ensure_array(key):
        if key not in d_dict:
            d_dict[key] = pikepdf.Array()
        return d_dict[key]

    ocgs = pdf.Root.OCProperties.get("/OCGs", [])
    target_ocgs = [ocg for ocg in ocgs if _get_obj_id(ocg) in target_ids]

    if not target_ocgs:
        return

    if action in ("show", "hide"):
        on_arr = _ensure_array("/ON")
        off_arr = _ensure_array("/OFF")

        # Remove from both arrays to prevent contradictory states
        _remove_targets_from_array(on_arr, target_ids)
        _remove_targets_from_array(off_arr, target_ids)

        target_arr = on_arr if action == "show" else off_arr
        target_arr.extend(target_ocgs)

    elif action in ("lock", "unlock"):
        locked_arr = _ensure_array("/Locked")
        _remove_targets_from_array(locked_arr, target_ids)

        if action == "lock":
            locked_arr.extend(target_ocgs)


def set_layer_usage(pdf, target_ids: set, action: str):
    """
    Applies usage overrides (print/noprint/screen/noscreen) to OCG dictionaries.
    """
    from pikepdf import Dictionary, Name

    if "/OCProperties" not in pdf.Root or "/OCGs" not in pdf.Root.OCProperties:
        return

    for ocg in pdf.Root.OCProperties.OCGs:
        _process_ocg_layer_usage(ocg, action, target_ids, Dictionary, Name)


def _process_ocg_layer_usage(ocg, action, target_ids, pikepdf_dictionary, pikepdf_name):
    if _get_obj_id(ocg) not in target_ids:
        return

    if "/Usage" not in ocg:
        ocg.Usage = pikepdf_dictionary()

    if action in ("print", "noprint"):
        if "/Print" not in ocg.Usage:
            ocg.Usage.Print = pikepdf_dictionary(
                Subtype=pikepdf_name.Print, PrintState=pikepdf_name.ON
            )
        ocg.Usage.Print.PrintState = pikepdf_name.ON if action == "print" else pikepdf_name.OFF

    elif action in ("screen", "noscreen"):
        if "/View" not in ocg.Usage:
            ocg.Usage.View = pikepdf_dictionary(ViewState=pikepdf_name.ON)
        ocg.Usage.View.ViewState = pikepdf_name.ON if action == "screen" else pikepdf_name.OFF
