vent_agent_2.0/extensions/main_fan/mine_catalog.py

"""批次固定首步：取得部门树，只保留正常生产煤矿，冻结可重读的矿井清单。"""
import asyncio
from collections import Counter
from copy import deepcopy
from datetime import datetime
import json
import re
from uuid import uuid4
from zoneinfo import ZoneInfo

import config
from extensions import get_service
from extensions.registry import load_extensions

DEPT_TREE_PATH = '/system/user/deptTree'
PRODUCTION_STATUS = '正常生产煤矿'


class MineCatalogError(ValueError):
    """上游目录结构或矿井身份不完整，不能形成后续批次输入。"""


def _identifier(value, field, *, allow_zero=False):
    if type(value) is int:
        value = str(value)
    if not isinstance(value, str) or not re.fullmatch(r'\d{1,64}', value):
        raise MineCatalogError(f'部门树 {field} 必须是完整的数字编码')
    if not allow_zero and int(value) == 0:
        raise MineCatalogError(f'部门树 {field} 不能为0')
    return value


def _name_and_status(label):
    # 只认末尾括号内的完整状态，避免把“非正常生产煤矿”等包含关系当成正常生产。
    for pattern in (r'(.+)（([^（）]+)）', r'(.+)\(([^()]+)\)'):
        match = re.fullmatch(pattern, label.strip())
        if match:
            return match[1].strip(), match[2].strip()
    return label.strip(), None


def parse_dept_tree(body):
    """纯代码遍历真实协议；保留全部煤矿供审计，mines 只包含精确匹配的叶节点。"""
    if (not isinstance(body, dict) or str(body.get('code')) not in {'0', '200'}
            or body.get('success') is False or not isinstance(body.get('data'), list)):
        raise MineCatalogError('部门树未返回成功的 data 节点数组')
    stack = [(node, []) for node in reversed(body['data'])]
    node_ids, mine_ids = set(), set()
    all_mines = []
    while stack:
        node, ancestors = stack.pop()
        if not isinstance(node, dict):
            raise MineCatalogError('部门树节点必须是对象')
        dept_id = _identifier(node.get('id'), 'id')
        parent_id = _identifier(node.get('parentId'), 'parentId', allow_zero=True)
        label = node.get('label')
        if not isinstance(label, str) or not label.strip():
            raise MineCatalogError(f'部门 {dept_id} 缺少有效 label')
        if dept_id in node_ids:
            raise MineCatalogError(f'部门树出现重复 id：{dept_id}')
        node_ids.add(dept_id)
        if ancestors and parent_id != ancestors[-1]['dept_id']:
            raise MineCatalogError(f'部门 {dept_id} 的 parentId 与树层级不一致')
        children = node.get('children')
        if children is not None and not isinstance(children, list):
            raise MineCatalogError(f'部门 {dept_id} 的 children 必须是数组或null')
        if children:
            path = [*ancestors, {'dept_id': dept_id, 'label': label}]
            stack.extend((child, path) for child in reversed(children))
            continue
        # 本接口约定叶节点为煤矿；缺身份时显式失败，不静默丢失煤矿或猜测 mineId。
        mine_id = node.get('mineId')
        if not isinstance(mine_id, str):
            raise MineCatalogError(f'煤矿部门 {dept_id} 缺少字符串 mineId')
        mine_id = _identifier(mine_id, 'mineId')
        if mine_id in mine_ids:
            raise MineCatalogError(f'部门树出现重复 mineId：{mine_id}')
        mine_ids.add(mine_id)
        mine_name, production_status = _name_and_status(label)
        if not mine_name:
            raise MineCatalogError(f'煤矿部门 {dept_id} 缺少名称')
        all_mines.append({
            'dept_id': dept_id, 'mine_id': mine_id, 'mine_name': mine_name,
            'production_status': production_status, 'label': label,
            'parent_dept_id': parent_id, 'department_path': deepcopy(ancestors),
            'raw_node': deepcopy(node),
        })
    mines = [mine for mine in all_mines if mine['production_status'] == PRODUCTION_STATUS]
    return {
        'counts': {'tree_nodes': len(node_ids), 'all_mines': len(all_mines),
                   'selected_mines': len(mines), 'excluded_mines': len(all_mines) - len(mines)},
        'status_counts': dict(Counter(mine['production_status'] or '未标注状态' for mine in all_mines)),
        'all_mines': all_mines, 'mines': mines,
    }


def _root():
    return config.ROOT / 'data/main_fan/mine_catalog'


def _write(path, value):
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(path.name + '.' + uuid4().hex + '.writing')
    try:
        temporary.write_text(json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False), encoding='utf-8')
        temporary.replace(path)
    finally:
        temporary.unlink(missing_ok=True)


def _archive(body, parsed, source):
    snapshot_id = uuid4().hex
    folder = _root() / snapshot_id
    snapshot = {
        'schema_version': 1, 'snapshot_id': snapshot_id, 'stage': 'mine_discovery',
        'status': 'completed', 'source': source,
        'filter': {'production_status': PRODUCTION_STATUS, 'match': 'exact_trailing_parentheses'},
        'counts': parsed['counts'], 'status_counts': parsed['status_counts'], 'mines': parsed['mines'],
    }
    _write(folder / 'request.json', source)
    _write(folder / 'dept-tree.json', body)
    _write(folder / 'all-mines.json', parsed['all_mines'])
    _write(folder / 'selected-mines.json', snapshot)
    # 只有成功取得、完整校验并归档的清单才更新 latest，失败不会覆盖上一次成功结果。
    _write(_root() / 'latest.json', {'snapshot_id': snapshot_id})
    return snapshot


async def refresh(*, gateway=None):
    """固定查询与筛选入口，不经过模型，也不在此步创建逐矿研判或推送任务。"""
    if not config.get_bool('SHANXI_GATEWAY_ENABLED', False):
        raise RuntimeError('矿井目录取数需要启用 SHANXI_GATEWAY_ENABLED')
    spec = load_extensions().business('main_fan')
    gateway = gateway if gateway is not None else get_service('shaanxi_gateway')
    body = await gateway.proxy(spec.gateway.agent_code, DEPT_TREE_PATH, params={})
    parsed = await asyncio.to_thread(parse_dept_tree, body)
    source = {'gateway_url': gateway.settings.base_url, 'agent_code': spec.gateway.agent_code,
              'method': 'GET', 'path': DEPT_TREE_PATH, 'params': {},
              'fetched_at': datetime.now(ZoneInfo('Asia/Shanghai')).isoformat()}
    return await asyncio.to_thread(_archive, body, parsed, source)


def load(snapshot_id=None):
    """重读冻结清单；后续逐矿任务应绑定 snapshot_id，不重新临时筛选。"""
    if snapshot_id is None:
        snapshot_id = json.loads((_root() / 'latest.json').read_text(encoding='utf-8'))['snapshot_id']
    if not isinstance(snapshot_id, str) or not re.fullmatch(r'[0-9a-f]{32}', snapshot_id):
        raise ValueError('无效的煤矿清单 snapshot_id')
    return json.loads((_root() / snapshot_id / 'selected-mines.json').read_text(encoding='utf-8'))