"""Divide una lista de AXNode en chunks de tamaño máximo para enviar a LLMs."""

import json


def chunk_ax_tree(nodes: list[dict], max_chars: int = 25000) -> list[list[dict]]:
    """Splittea AX tree en chunks de hasta max_chars cuando se serializa a JSON.

    Estrategia DFS desde el nodo raíz:
      - Acumula nodos en el chunk actual.
      - Cuando añadir el siguiente subárbol superaría max_chars, cierra el chunk
        y abre uno nuevo.
      - Cada chunk incluye al inicio los nodos "context" (path desde root hasta
        el primer nodo del chunk) para que el LLM tenga orientación estructural.

    Args:
        nodes: Lista de AXNode en formato CDP (con nodeId, childIds, role, name).
        max_chars: Límite de caracteres al serializar el chunk a JSON. Default 25000.

    Returns:
        Lista de chunks. Cada chunk es una lista de AXNode. El primer elemento de
        cada chunk (salvo el primero) es un nodo sintético con role="context" que
        contiene el path desde la raíz.
    """
    if not nodes:
        return []

    # Construir lookup y encontrar raíz
    by_id: dict[str, dict] = {n["nodeId"]: n for n in nodes}
    all_children: set[str] = set()
    for n in nodes:
        for cid in n.get("childIds", []):
            all_children.add(cid)
    roots = [n for n in nodes if n["nodeId"] not in all_children]
    if not roots:
        # Fallback: usar el primer nodo como raíz
        roots = [nodes[0]]

    # DFS para obtener orden de visita + path desde raíz
    visit_order: list[tuple[dict, list[dict]]] = []  # (nodo, ancestors)

    def dfs(node: dict, ancestors: list[dict]) -> None:
        visit_order.append((node, list(ancestors)))
        for cid in node.get("childIds", []):
            child = by_id.get(cid)
            if child:
                dfs(child, ancestors + [node])

    for root in roots:
        dfs(root, [])

    # Agrupar en chunks respetando max_chars
    chunks: list[list[dict]] = []
    current_chunk: list[dict] = []
    current_chars: int = 2  # "[]"

    for node, ancestors in visit_order:
        node_json = json.dumps(node, ensure_ascii=False)
        node_chars = len(node_json) + 2  # ", " separator

        if current_chunk and current_chars + node_chars > max_chars:
            chunks.append(current_chunk)
            # Nuevo chunk con contexto (path desde root)
            current_chunk = []
            current_chars = 2
            if ancestors:
                context_node = {
                    "nodeId": "context",
                    "role": {"value": "context"},
                    "name": {"value": "path-from-root"},
                    "childIds": [],
                    "context_path": [
                        {
                            "nodeId": a["nodeId"],
                            "role": a.get("role", {}).get("value", "") if isinstance(a.get("role"), dict) else str(a.get("role", "")),
                            "name": a.get("name", {}).get("value", "") if isinstance(a.get("name"), dict) else str(a.get("name", "")),
                        }
                        for a in ancestors
                    ],
                }
                ctx_chars = len(json.dumps(context_node, ensure_ascii=False)) + 2
                current_chunk.append(context_node)
                current_chars += ctx_chars

        current_chunk.append(node)
        current_chars += node_chars

    if current_chunk:
        chunks.append(current_chunk)

    return chunks
