from sqlalchemy.orm import noload
from sqlalchemy.sql._typing import ColumnExpressionArgument
from typing import List, Dict, Any, Optional
from iti.applications.extensions import db
from iti.applications.models import SysMenu
from iti.applications.common.enums import StatusEnum, MenuTypeEnum
from iti.applications.common.utils import (
    build_tree_from_list,
)
from sqlalchemy import select


def get_menu_tree(
    parent_id: Optional[str] = None,
    include_disabled: bool = True,
    type_filter: Optional[List[MenuTypeEnum]] = None,
    status_filter: Optional[List[StatusEnum]] = None,
) -> List[Dict[str, Any]]:
    """
    获取菜单树
    Args:
        parent_id: 父菜单ID 不填为所有树，若填写则只返回该父菜单的子菜单树(包括父菜单)
        include_disabled: 是否包含禁用菜单
        type_filter: 菜单类型过滤 不填为不限制
        status_filter: 状态过滤 不填为不限制
    """
    # 获取所有菜单数据
    cte_query = build_descendants_cte(SysMenu.parent_id == parent_id)
    if not include_disabled:
        cte_query = cte_query.filter(SysMenu.status == StatusEnum.ENABLED)
    if type_filter:
        cte_query = cte_query.filter(SysMenu.type.in_(type_filter))
    if status_filter:
        cte_query = cte_query.filter(SysMenu.status.in_(status_filter))
    cte_query = cte_query.order_by(SysMenu.parent_id, SysMenu.sort).options(
        noload(SysMenu.parent), noload(SysMenu.children)
    )
    # 使用增强的树工具构建树结构
    return build_tree_from_list(db.session.scalars(cte_query).all())


def build_descendants_cte(*criterion: ColumnExpressionArgument[bool]):
    """
    构建获取后代的递归CTE
    """
    hierarchy = (
        select(SysMenu.id, SysMenu.parent_id)
        .filter(*criterion)
        .cte("hierarchy", recursive=True)
    )
    bt = select(SysMenu.id, SysMenu.parent_id).join(
        hierarchy, hierarchy.c.id == SysMenu.parent_id
    )
    q = hierarchy.union_all(bt)

    return select(SysMenu).join(q, SysMenu.id == q.c.id)


def build_ancestors_cte(*criterion: ColumnExpressionArgument[bool]):
    """
    构建获取祖先的递归CTE
    """
    hierarchy = (
        select(SysMenu.id, SysMenu.parent_id)
        .filter(*criterion)
        .cte("hierarchy", recursive=True)
    )
    bt = select(SysMenu.id, SysMenu.parent_id).join(
        hierarchy, hierarchy.c.parent_id == SysMenu.id
    )
    q = hierarchy.union_all(bt)

    return select(SysMenu).join(q, SysMenu.id == q.c.id)