from typing import List, Dict, Tuple, Any, Optional
from collections import defaultdict
from ..schemas import OverallState, GroupTask, FieldSpec


def flatten_leaf_tasks(specs: List[FieldSpec], base_path: Optional[List[str]] = None) -> List[Tuple[Tuple[str, ...], str]]:
    """
    返回：[(path_tuple, leaf_key), ...]
    """
    base_path = base_path or []
    out: List[Tuple[Tuple[str, ...], str]] = []

    def _as_list(x: Any) -> List[Any]:
        if x is None:
            return []
        if isinstance(x, list):
            return x
        return [x]

    for item in specs or []:
        if isinstance(item, str):
            out.append((tuple(base_path), item))
            continue

        if isinstance(item, dict):
            for parent, children in item.items():
                for child in _as_list(children):
                    if isinstance(child, str):
                        out.append((tuple(base_path + [parent]), child))
                    elif isinstance(child, dict):
                        out.extend(flatten_leaf_tasks([child], base_path + [parent]))
                    else:
                        pass
            continue

    return out


def create_plan_node():
    """创建规划节点：将输入的字段模板规划为多个独立的抽取任务"""
    def plan_node(state: OverallState):
        leaf = flatten_leaf_tasks(state["fields"])
        grouped: Dict[Tuple[str, ...], List[str]] = defaultdict(list)
        for path, key in leaf:
            grouped[path].append(key)

        tasks: List[GroupTask] = []
        for path, keys in grouped.items():
            # 去重保持顺序
            seen = set()
            uniq = []
            for k in keys:
                if k not in seen:
                    seen.add(k)
                    uniq.append(k)
            tasks.append(GroupTask(path=path, keys=tuple(uniq)))

        tasks.sort(key=lambda t: (len(t.path), t.path))
        return {"tasks": tasks}
    
    return plan_node

