"""Shared layer-tree grouping helper.

All AI Edit outputs (generated rasters, Mark up annotations, Vectorize
polygons) live under a single group at the top of the layer tree so
users keep their existing layer organisation untouched.

Generations land as flat children of the AI-Edit group. A per-raster
sub-group is created lazily the first time a raster is vectorized: at
that point the raster is moved into a new sub-group named after it and
its vector layer joins it. Subsequent vectorizations of the same raster
land in the same sub-group. Generations that are never vectorized stay
flat, so the tree only grows nesting when bundling is actually useful.
"""
from __future__ import annotations

from typing import Callable

from qgis.core import QgsLayerTreeGroup, QgsLayerTreeLayer, QgsMapLayer, QgsProject

AI_EDIT_GROUP_NAME = "AI-Edit"
# Custom property stamped on the group we own. Lets us distinguish our
# group from a user-created or other-plugin "AI-Edit" group that happens
# to share the name, so we never absorb someone else's layers.
_OWNERSHIP_PROPERTY = "terralab/ai_edit_group"
# Per-generation sub-group marker + slot for the source raster layer id
# so vector layers know which sub-group to join even after a project reload.
_GENERATION_SUBGROUP_PROPERTY = "terralab/ai_edit_generation"
_GENERATION_SOURCE_LAYER_PROPERTY = "terralab/ai_edit_source_layer_id"
# Custom property stamped on the shared Mark up annotation layer. The
# layer is unique per project and lives at the bottom of the AI-Edit
# group so new generations stack above it without pushing it around.
MARKUP_LAYER_PROPERTY = "ai_edit_markup"


def _walk_groups(node):
    """Yield every QgsLayerTreeGroup descendant of ``node`` (depth-first)."""
    for child in node.children():
        if isinstance(child, QgsLayerTreeGroup):
            yield child
            yield from _walk_groups(child)


def find_ai_edit_group() -> QgsLayerTreeGroup | None:
    """The AI-Edit group anywhere in the tree, or None. Never creates it.

    Marker first, name second, mirroring ``get_or_create_ai_edit_group``'s
    resolution order minus the creating step - custom properties on tree nodes
    do not always survive a project save/reload in QGIS 3.x, so the name is the
    fallback. The walk is recursive because the user may have dragged the group
    into a folder of their own.
    """
    root = QgsProject.instance().layerTreeRoot()
    for child in _walk_groups(root):
        if child.customProperty(_OWNERSHIP_PROPERTY):
            return child
    for child in _walk_groups(root):
        if child.name() == AI_EDIT_GROUP_NAME:
            return child
    return None


def get_or_create_ai_edit_group() -> QgsLayerTreeGroup:
    """Return the AI-Edit group, creating it at the top of the tree if absent.

    Resolution order:
      1. Any group anywhere in the tree that carries our ownership marker
         (lets the user drag the AI-Edit group into a folder of their own
         layout without us spawning a duplicate next time).
      2. Any group anywhere in the tree named ``AI-Edit`` (custom
         properties on tree nodes do not always survive project save/reload
         in QGIS 3.x, so we fall back to name matching and re-stamp the
         marker so future calls hit the fast path).
      3. Create a fresh group at the top of the tree.

    The recursive walk is the fix for the audit finding "user drags AI-Edit
    into a sub-folder, next reload creates a second group at root".
    """
    root = QgsProject.instance().layerTreeRoot()
    for child in _walk_groups(root):
        if child.customProperty(_OWNERSHIP_PROPERTY):
            return child
    for child in _walk_groups(root):
        if child.name() == AI_EDIT_GROUP_NAME:
            child.setCustomProperty(_OWNERSHIP_PROPERTY, True)
            return child
    group = root.insertGroup(0, AI_EDIT_GROUP_NAME)
    group.setCustomProperty(_OWNERSHIP_PROPERTY, True)
    # Creating the group at index 0 pushes the Mark up layer down; bounce it
    # back to the very top so its annotations stay visible above the group.
    pin_markup_to_top()
    return group


def collect_ai_edit_layer_ids() -> set[str]:
    """Layer IDs of everything under the AI-Edit group (generated rasters,
    vectorize outputs), gathered recursively. Note the Mark up annotation layer
    is NOT included: it lives at the layer-tree root, not in this group, and is
    handled separately by the export's markup_layer path.

    Does NOT create the group: returns an empty set when it is absent. Used to
    render the "original" (pre-AI-edit) base for a retry - the caller drops
    these layers from the export so the model sees the clean map.
    """
    group = find_ai_edit_group()
    if group is None:
        return set()

    ids: set[str] = set()

    def _collect(node) -> None:
        for child in node.children():
            if isinstance(child, QgsLayerTreeGroup):
                _collect(child)
            elif isinstance(child, QgsLayerTreeLayer):
                layer_id = child.layerId()
                if layer_id:
                    ids.add(layer_id)

    _collect(group)
    return ids


def set_ai_edit_layers_checked(checked: bool, except_ids: set[str] | None = None) -> None:
    """Check or uncheck the visibility box of every layer under the AI-Edit group.

    Drives the post-generation base choice on the canvas: picking "original"
    unchecks the AI result layers so the map shows the clean base that will
    actually be sent; picking the AI-edited base re-checks them so the layer the
    user is about to edit stays visible. ``except_ids`` is left untouched - the
    active Mark up layer is user guidance, not an AI edit, so the caller passes
    it here. No-op when the group is absent.
    """
    group = find_ai_edit_group()
    if group is None:
        return
    skip = except_ids or set()

    def _apply(node) -> None:
        for child in node.children():
            if isinstance(child, QgsLayerTreeGroup):
                _apply(child)
            elif isinstance(child, QgsLayerTreeLayer):
                if child.layerId() and child.layerId() not in skip:
                    child.setItemVisibilityChecked(checked)

    _apply(group)


def find_generation_subgroup_for_layer(layer_id: str) -> QgsLayerTreeGroup | None:
    """Return the per-raster sub-group that contains ``layer_id``.

    Walks the AI-Edit group's sub-groups and looks for a layer node whose
    id matches. Used by Vectorize to drop vector outputs into the same
    sub-group as their source raster instead of a flat group.
    """
    parent = get_or_create_ai_edit_group()
    for child in parent.children():
        if not isinstance(child, QgsLayerTreeGroup):
            continue
        for sub in child.children():
            if isinstance(sub, QgsLayerTreeLayer) and sub.layerId() == layer_id:
                return child
    return None


def promote_layer_to_own_subgroup(layer_id: str) -> QgsLayerTreeGroup | None:
    """Move a flat AI-Edit child raster into a fresh sub-group named after it.

    Called the first time a raster is vectorized so the raster + its
    vector(s) get bundled together. If the raster is already inside a
    sub-group, returns that existing sub-group. Returns None if the
    layer can't be located under the AI-Edit group.
    """
    parent = get_or_create_ai_edit_group()
    project = QgsProject.instance()

    layer = project.mapLayer(layer_id)
    if layer is None:
        return None

    source_node = None
    for child in parent.children():
        if isinstance(child, QgsLayerTreeLayer) and child.layerId() == layer_id:
            source_node = child
            break
    if source_node is None:
        return find_generation_subgroup_for_layer(layer_id)

    base_name = layer.name()
    existing = {
        child.name()
        for child in parent.children()
        if isinstance(child, QgsLayerTreeGroup)
    }
    final_name = base_name
    counter = 2
    while final_name in existing:
        final_name = f"{base_name} ({counter})"
        counter += 1

    subgroup = parent.insertGroup(0, final_name)
    subgroup.setCustomProperty(_GENERATION_SUBGROUP_PROPERTY, True)
    subgroup.setExpanded(True)

    # Attach the clone BEFORE removing the original. QgsLayerTreeRegistryBridge
    # auto-removes a layer from the project when its last tree node disappears,
    # so the layer must keep at least one tree reference at all times during
    # the move. Detaching then reattaching deletes the raster.
    clone = source_node.clone()
    subgroup.addChildNode(clone)
    parent.removeChildNode(source_node)
    return subgroup


def add_layer_to_ai_edit_top(layer: QgsMapLayer) -> QgsLayerTreeLayer:
    """Insert ``layer`` at the top of the AI-Edit group so the newest output
    is always the topmost child. The Mark up annotation layer lives at the
    tree root above the group (not inside it), so new outputs go straight to
    index 0 without disturbing it.
    """
    group = get_or_create_ai_edit_group()
    return group.insertLayer(0, layer)


def bring_ai_edit_group_to_front() -> bool:
    """Move the AI-Edit group back to the top of the tree root.

    The group is allowed to live anywhere (``find_ai_edit_group`` looks
    everywhere, so the user can file it in a folder of their own). The cost is
    that a group parked under an opaque basemap draws every generation BEHIND
    the imagery it edited: the raster is written, the credit is spent, and the
    canvas does not change. This is the one-click repair offered when that is
    detected; it is never applied on its own, since the layout is the user's.

    Returns whether the group ends up at the front. Idempotent.
    """
    root = QgsProject.instance().layerTreeRoot()
    group = find_ai_edit_group()
    if group is None:
        return False
    parent = group.parent()
    if parent is root and root.children() and root.children()[0] is group:
        return True
    # Clone-then-remove keeps a tree reference alive throughout, otherwise
    # QgsLayerTreeRegistryBridge drops the member layers from the project
    # (same gotcha as pin_markup_to_top and promote_layer_to_own_subgroup).
    clone = group.clone()
    root.insertChildNode(0, clone)
    if parent is not None:
        parent.removeChildNode(group)
    pin_markup_to_top()
    return True


def pin_markup_to_top() -> None:
    """Keep the Mark up layer at the very top of the tree, above the AI-Edit
    group, so its annotations always render over everything else.

    Idempotent. The Mark up layer lives as a direct child of the tree root
    (not inside the AI-Edit group) and is found by ``MARKUP_LAYER_PROPERTY``
    so it survives rename + reload.
    """
    root = QgsProject.instance().layerTreeRoot()
    children = list(root.children())
    markup_node = None
    for child in children:
        if isinstance(child, QgsLayerTreeLayer):
            layer = child.layer()
            if layer is not None and layer.customProperty(MARKUP_LAYER_PROPERTY):
                markup_node = child
                break
    if markup_node is None or children.index(markup_node) == 0:
        return
    # Clone-then-remove keeps at least one tree reference alive at all
    # times, otherwise QgsLayerTreeRegistryBridge would wipe the layer
    # from the project (same gotcha as promote_layer_to_own_subgroup).
    clone = markup_node.clone()
    root.insertChildNode(0, clone)
    root.removeChildNode(markup_node)


def most_recent_ai_edit_output(
    predicate: Callable[[QgsMapLayer], bool] | None = None,
) -> QgsMapLayer | None:
    """Return the topmost layer under the AI-Edit group passing ``predicate``.

    New generations are inserted at index 0 so the topmost match is the
    most recent. Walks per-generation sub-groups one level deep so vector
    outputs nested under their source raster are reachable too.
    """
    group = find_ai_edit_group()
    if group is None:
        return None
    for sub in group.children():
        if isinstance(sub, QgsLayerTreeLayer):
            layer = sub.layer()
            if layer is not None and (predicate is None or predicate(layer)):
                return layer
        elif isinstance(sub, QgsLayerTreeGroup):
            for leaf in sub.children():
                if isinstance(leaf, QgsLayerTreeLayer):
                    layer = leaf.layer()
                    if layer is not None and (predicate is None or predicate(layer)):
                        return layer
    return None


def pick_default_layer(
    predicate: Callable[[QgsMapLayer], bool],
) -> QgsMapLayer | None:
    """Cascade picker shared by tool panels for their "From" combo.

    Priority: 1) the most recent AI-Edit output that passes ``predicate``;
    2) the user's currently active layer in the QGIS Layers panel if it
    passes; 3) None (caller falls back to its own logic).
    """
    most_recent = most_recent_ai_edit_output(predicate)
    if most_recent is not None:
        return most_recent
    try:
        from qgis.utils import iface as _iface
        active = _iface.activeLayer() if _iface is not None else None
    except Exception:
        active = None
    if active is not None:
        try:
            if predicate(active):
                return active
        except Exception:  # nosec B110
            pass
    return None


def drop_from_snapping(layer: QgsMapLayer) -> None:
    """Scrub a scratch layer from the project snapping config, right AFTER it
    is registered (addMapLayer auto-enrolls vector layers there).

    QgsSnappingConfig keeps raw layer pointers: a freed working layer leaves a
    dangling entry that crashes the NEXT project save, typically the
    save-on-exit (upstream qgis/QGIS#37505, #42651). A layer with no entry can
    never dangle, and scratch layers have no business being snap targets.
    """
    try:
        project = QgsProject.instance()
        cfg = project.snappingConfig()
        cfg.removeLayers([layer])
        project.setSnappingConfig(cfg)
    except Exception:  # nosec B110
        pass
