KCGL/inventory-backend/app/api/v1/warehouse.py

# inventory-backend/app/api/v1/warehouse.py
from flask import Blueprint, request, jsonify
from flask_jwt_extended import jwt_required
from app.extensions import db
from app.models.system import SysWarehouseLocation

warehouse_bp = Blueprint('warehouse', __name__, url_prefix='/api/v1/warehouse')


def build_tree(nodes, parent_id=None):
    """
    将平铺的数据构建为树形结构（O(N) 内存组装，避免递归时每层全量扫描导致 O(N²)）

    做法：先把全部节点按 parent_id 分组到 Map，再从根出发逐层用 Map 取子节点组装。
    每个节点只被处理一次，显著快于"每次递归 for 遍历全表"的旧实现。
    """
    by_parent = {}
    for node in nodes:
        by_parent.setdefault(node.parent_id, []).append(node)

    def assemble(pid):
        kids = sorted(by_parent.get(pid, []), key=lambda x: (x.name or ''))
        out = []
        for k in kids:
            d = k.to_dict()
            d['children'] = assemble(k.id)
            out.append(d)
        return out

    return assemble(parent_id)


@warehouse_bp.route('/tree', methods=['GET'])
def get_tree():
    """
    获取库位树形结构

    查询参数:
      prefixes —— 可选，逗号分隔的顶层前缀，例如 ?prefixes=Y 或 ?prefixes=C,L
                  只返回**顶层** name / full_path 命中这些前缀的根节点及其完整子树；
                  不传则返回全量。

    用途：前端按公司精简拉取（IRIS 只要 Y*，LICA 只要 C*/L*），
    在**保留完整子树**的前提下减少节点数与传输量 —— 不能退回懒加载，
    因为 setCheckedKeys / getCheckedNodes 依赖全树已构建。
    """
    try:
        raw_prefixes = request.args.get('prefixes', '', type=str)
        prefixes = [p.strip().upper() for p in raw_prefixes.split(',') if p.strip()]

        # 查询所有库位，按 name 升序排序
        all_locations = SysWarehouseLocation.query.order_by(SysWarehouseLocation.name.asc()).all()

        # 构建树形结构（O(N) 内存组装，见 build_tree）
        tree_data = build_tree(all_locations, parent_id=None)

        # ★ 只在**顶层**做前缀过滤：命中即整棵子树保留，不递归裁剪，
        #   避免把子树打散导致前端勾选语义错乱
        if prefixes:
            def _hit(node):
                name = str(node.get('name') or '').upper()
                path = str(node.get('full_path') or '').upper()
                return any(name.startswith(p) or path.startswith(p) for p in prefixes)

            tree_data = [n for n in tree_data if _hit(n)]

        return jsonify({
            'code': 200,
            'msg': 'success',
            'data': tree_data
        })
    except Exception as e:
        return jsonify({
            'code': 500,
            'msg': str(e),
            'data': None
        }), 500


@warehouse_bp.route('/children', methods=['GET'])
def get_children():
    """
    懒加载：获取指定库位的直接子节点（parent_id 省略/为空 = 顶层）。
    每个节点附带 has_children 标记，前端据此渲染「进入下级」而非点进去才知道。
    与 /tree 行为一致（不额外过滤 is_enabled、按 name 升序）。
    """
    try:
        parent_id = request.args.get('parent_id', type=int)
        if parent_id is None:
            nodes = SysWarehouseLocation.query.filter(
                SysWarehouseLocation.parent_id.is_(None)
            ).order_by(SysWarehouseLocation.name.asc()).all()
        else:
            nodes = SysWarehouseLocation.query.filter(
                SysWarehouseLocation.parent_id == parent_id
            ).order_by(SysWarehouseLocation.name.asc()).all()

        # 一次查询所有"有子节点"的 parent_id，用于 has_children 判断（避免 N+1）
        parent_with_children = set(
            cid for (cid,) in db.session.query(SysWarehouseLocation.parent_id)
            .filter(SysWarehouseLocation.parent_id.isnot(None)).distinct().all()
        )

        data = []
        for node in nodes:
            d = node.to_dict()
            d['has_children'] = node.id in parent_with_children
            data.append(d)

        return jsonify({'code': 200, 'msg': 'success', 'data': data})
    except Exception as e:
        return jsonify({'code': 500, 'msg': str(e), 'data': None}), 500


@warehouse_bp.route('', methods=['POST'])
@jwt_required()
def create_location():
    """
    创建库位
    """
    try:
        data = request.get_json()
        name = data.get('name', '').strip()
        parent_id = data.get('parent_id')  # None 表示顶级
        is_enabled = data.get('is_enabled', True)
        
        if not name:
            return jsonify({'code': 400, 'msg': '库位名称不能为空', 'data': None})
        
        # 计算 level 和 full_path
        if parent_id is None:
            level = 0
            full_path = name
            parent_full_path = ''
        else:
            parent = SysWarehouseLocation.query.get(parent_id)
            if not parent:
                return jsonify({'code': 400, 'msg': '父级库位不存在', 'data': None})
            level = parent.level + 1
            parent_full_path = parent.full_path or ''
            full_path = f"{parent_full_path}/{name}" if parent_full_path else name
        
        location = SysWarehouseLocation(
            name=name,
            parent_id=parent_id,
            full_path=full_path,
            level=level,
            is_enabled=is_enabled
        )
        db.session.add(location)
        db.session.commit()
        
        return jsonify({
            'code': 200,
            'msg': '创建成功',
            'data': location.to_dict()
        })
    except Exception as e:
        db.session.rollback()
        return jsonify({
            'code': 500,
            'msg': str(e),
            'data': None
        }), 500


@warehouse_bp.route('/<int:location_id>', methods=['PUT'])
@jwt_required()
def update_location(location_id):
    """
    更新库位
    """
    try:
        data = request.get_json()
        location = SysWarehouseLocation.query.get(location_id)
        
        if not location:
            return jsonify({'code': 404, 'msg': '库位不存在', 'data': None})
        
        # 更新名称
        if 'name' in data and data['name']:
            new_name = data['name'].strip()
            if new_name != location.name:
                # 需要更新 full_path
                parent = location.parent
                if parent:
                    location.full_path = f"{parent.full_path}/{new_name}" if parent.full_path else new_name
                else:
                    location.full_path = new_name
                location.name = new_name
        
        # 更新启用状态
        if 'is_enabled' in data:
            location.is_enabled = data['is_enabled']
        
        db.session.commit()
        
        return jsonify({
            'code': 200,
            'msg': '更新成功',
            'data': location.to_dict()
        })
    except Exception as e:
        db.session.rollback()
        return jsonify({
            'code': 500,
            'msg': str(e),
            'data': None
        }), 500


@warehouse_bp.route('/<int:location_id>', methods=['DELETE'])
@jwt_required()
def delete_location(location_id):
    """
    删除库位（级联删除子库位）
    """
    try:
        location = SysWarehouseLocation.query.get(location_id)

        if not location:
            return jsonify({'code': 404, 'msg': '库位不存在', 'data': None})

        # 在删除前提取属性，避免 commit 后访问已删除对象
        deleted_loc_name = location.name

        # 递归删除所有子库位
        def delete_recursive(loc):
            # 先删除所有子节点
            children = SysWarehouseLocation.query.filter_by(parent_id=loc.id).all()
            for child in children:
                delete_recursive(child)
            # 再删除自身
            db.session.delete(loc)

        delete_recursive(location)
        db.session.commit()

        return jsonify({
            'code': 200,
            'msg': '删除成功',
            'deleted_location': deleted_loc_name
        })
    except Exception as e:
        db.session.rollback()
        return jsonify({
            'code': 500,
            'msg': str(e),
            'data': None
        }), 500


@warehouse_bp.route('/batch', methods=['DELETE'])
@jwt_required()
def batch_delete_locations():
    """
    批量删除库位
    """
    try:
        ids = request.get_json()
        if not ids or not isinstance(ids, list):
            return jsonify({'code': 400, 'msg': '请提供要删除的库位ID列表', 'data': None})

        deleted_count = 0
        deleted_names = []

        for loc_id in ids:
            location = SysWarehouseLocation.query.get(loc_id)
            if not location:
                continue

            # 在删除前提取属性
            deleted_names.append(location.name)

            # 递归删除
            def delete_recursive(loc):
                children = SysWarehouseLocation.query.filter_by(parent_id=loc.id).all()
                for child in children:
                    delete_recursive(child)
                db.session.delete(loc)

            delete_recursive(location)
            deleted_count += 1

        db.session.commit()

        return jsonify({
            'code': 200,
            'msg': f'删除成功，共删除 {deleted_count} 个库位',
            'data': {'deleted_count': deleted_count, 'deleted_names': deleted_names}
        })
    except Exception as e:
        db.session.rollback()
        return jsonify({
            'code': 500,
            'msg': str(e),
            'data': None
        }), 500


@warehouse_bp.route('/batch-generate', methods=['POST'])
@jwt_required()
def batch_generate_locations():
    """
    规则化批量新增库位
    """
    MAX_TOTAL = 3000  # 单次最多生成数量限制
    try:
        data = request.get_json()
        parent_id = data.get('parent_id')
        rules = data.get('rules', [])

        if not rules:
            return jsonify({'code': 400, 'msg': '请提供生成规则', 'data': None})

        # 验证规则并计算总数
        total_count = 1
        for rule in rules:
            start = rule.get('start', 1)
            end = rule.get('end', 1)
            total_count *= max(0, end - start + 1)

        if total_count > MAX_TOTAL:
            return jsonify({'code': 400, 'msg': f'单次生成数量不能超过 {MAX_TOTAL} 个，当前计划生成 {total_count} 个', 'data': None})

        # 初始化父级列表
        if parent_id:
            parent = SysWarehouseLocation.query.get(parent_id)
            if not parent:
                return jsonify({'code': 404, 'msg': '父级库位不存在', 'data': None})
            current_parents = [parent_id]
        else:
            current_parents = [None]

        # 逐层处理规则
        generated_ids = []

        for rule in rules:
            prefix = rule.get('prefix', '')
            start = rule.get('start', 1)
            end = rule.get('end', 1)
            pad = rule.get('pad', 1)

            new_locations = []

            for parent_id in current_parents:
                # 1. 动态获取当前特定父节点的信息（严禁放循环外面共享！）
                if parent_id is None:
                    current_level = 0
                    current_parent_path = ''
                else:
                    p = SysWarehouseLocation.query.get(parent_id)
                    current_level = (p.level + 1) if p else 0
                    current_parent_path = p.full_path if p and p.full_path else ''

                # 2. 生成当前父节点下的专属子节点
                for num in range(start, end + 1):
                    name = f"{prefix}{str(num).zfill(pad)}"
                    # 路径由当前特定的 current_parent_path 决定
                    full_path = f"{current_parent_path}/{name}" if current_parent_path else name

                    location = SysWarehouseLocation(
                        name=name,
                        parent_id=parent_id,
                        full_path=full_path,
                        level=current_level,
                        is_enabled=True
                    )
                    db.session.add(location)
                    new_locations.append(location)

            # 单层循环结束后再 flush 和获取新 ID 列表
            db.session.flush()
            current_parents = [loc.id for loc in new_locations]
            generated_ids.extend(current_parents)

        db.session.commit()

        return jsonify({
            'code': 200,
            'msg': f'生成成功，共生成 {len(generated_ids)} 个库位',
            'data': {'generated_count': len(generated_ids), 'generated_ids': generated_ids}
        })
    except Exception as e:
        db.session.rollback()
        return jsonify({
            'code': 500,
            'msg': str(e),
            'data': None
        }), 500