vent_agent_2.0/extensions/main_fan/crawler_reports.py

"""资料采集 API + 只读挂载目录；原件复制进 workspace 后才交给 Agent。"""
import asyncio
from contextlib import asynccontextmanager
from dataclasses import asdict, dataclass
from datetime import datetime
import hashlib
import json
import math
import os
from pathlib import Path
import re
import time
from urllib.parse import quote, urlsplit

import httpx

import config
from extensions.main_fan.mine_catalog import _write
from extensions.main_fan.report_api import ReportClient, mock_report_client
from obs import get_logger

logger = get_logger(__name__)


def _part(value):
    if not isinstance(value, str) or not value.strip() or value in {'.', '..'} or any(c in value for c in '/\\\x00'):
        raise ValueError('资料目录中的矿名或类型名称不是有效的单级目录名')
    return value


def _check_success(body):
    if (not isinstance(body, dict) or body.get('success') is False
            or ('code' in body and str(body['code']) not in {'0', '200'})):
        raise ValueError('资料服务未返回成功的 JSON 对象')


def _data(body):
    _check_success(body)
    data = body.get('data', body)
    if not isinstance(data, dict):
        raise ValueError('资料服务 data 必须是对象')
    return data


@dataclass(frozen=True)
class CrawlerSettings:
    api_url: str
    source_dir: str
    report_type: str
    period: str = ''
    timeout: float = 30
    poll_seconds: float = 3
    task_timeout: float = 900
    max_files: int = 30
    max_file_mb: int = 50

    @classmethod
    def from_config(cls):
        value = cls(config.get_str('MAIN_FAN_REPORT_API_URL', 'http://10.244.3.155:19010').rstrip('/'),
            str(Path(config.get_str('MAIN_FAN_REPORT_SOURCE_DIR', '/data/source_data')).absolute()),
            config.get_str('MAIN_FAN_REPORT_TYPE', '通风机性能测定报告'),
            config.get_str('MAIN_FAN_REPORT_PERIOD'),
            config.get_float('MAIN_FAN_REPORT_TIMEOUT_SECONDS', 30),
            config.get_float('MAIN_FAN_REPORT_POLL_SECONDS', 3),
            config.get_float('MAIN_FAN_REPORT_TASK_TIMEOUT_SECONDS', 900),
            config.get_int('MAIN_FAN_REPORT_MAX_FILES', 30),
            config.get_int('MAIN_FAN_REPORT_MAX_FILE_MB', 50))
        parts = urlsplit(value.api_url)
        if (parts.scheme not in {'http', 'https'} or not parts.hostname or parts.username
                or parts.password or parts.query or parts.fragment):
            raise ValueError('MAIN_FAN_REPORT_API_URL 必须是无内嵌凭证的 HTTP(S) 地址')
        _part(value.report_type)
        if value.period and not re.fullmatch(r'\d{4}-(0[1-9]|1[0-2])', value.period):
            raise ValueError('MAIN_FAN_REPORT_PERIOD 必须为 YYYY-MM，留空使用批次截止月份')
        if (any(not math.isfinite(n) or n <= 0 for n in (value.timeout, value.poll_seconds, value.task_timeout))
                or value.poll_seconds > 60 or not 1 <= value.max_files <= 200 or not 1 <= value.max_file_mb <= 500):
            raise ValueError('资料服务超时、轮询间隔或文件上限配置无效')
        return value


def report_source(pdf_path=None):
    if pdf_path is not None:
        path = Path(pdf_path).resolve()
        if not path.is_file() or path.suffix.lower() != '.pdf':
            raise ValueError('模拟报告接口需要已有的测试 PDF')
        with path.open('rb') as stream:
            sha = hashlib.file_digest(stream, 'sha256').hexdigest()
        return {'mode': 'mock', 'sha256': sha}
    mode = config.get_str('MAIN_FAN_REPORT_MODE', 'crawler').lower()
    if mode != 'crawler':
        raise ValueError('mock 模式需配置 MAIN_FAN_MOCK_REPORT_PDF；正式模式应为 crawler')
    return {'mode': 'crawler', **asdict(CrawlerSettings.from_config())}


@asynccontextmanager
async def open_report_source(source, pdf_path=None):
    if source['mode'] == 'mock':
        async with mock_report_client(pdf_path) as http:
            yield ReportClient(http)
    else:
        settings = CrawlerSettings(**{key: value for key, value in source.items() if key != 'mode'})
        async with httpx.AsyncClient(base_url=settings.api_url + '/', timeout=settings.timeout,
                trust_env=False, follow_redirects=False) as http:
            yield CrawlerReportClient(http, settings)


class CrawlerReportClient:
    def __init__(self, http, settings):
        self.http, self.settings = http, settings
        self.cache = {}
        self.lock = asyncio.Lock()

    async def _request(self, method, path, **kwargs):
        started = time.monotonic()
        try:
            response = await self.http.request(method, path.lstrip('/'), **kwargs)
        except httpx.RequestError as exc:
            logger.warning('资料服务连接失败', extra={'kv': {'event': 'report.http.failed',
                'method': method, 'path': path, 'error_type': type(exc).__name__}})
            raise ValueError(f'资料服务 {method} {path} 连接失败：{type(exc).__name__}') from None
        logger.debug('资料服务请求', extra={'kv': {'event': 'report.http', 'method': method,
            'path': path, 'http_status': response.status_code,
            'duration_ms': round((time.monotonic() - started) * 1000)}})
        if response.status_code != 409 and not 200 <= response.status_code < 300:
            raise ValueError(f'资料服务 {method} {path} 返回 HTTP {response.status_code}')
        try:
            body = response.json()
        except ValueError:
            raise ValueError(f'资料服务 {method} {path} 未返回 JSON') from None
        return response.status_code, body

    async def list_for_mine(self, mine, cutoff, receipt_dir):
        period = self.settings.period or datetime.fromisoformat(cutoff).strftime('%Y-%m')
        key = (str(mine['dept_id']), mine['mine_name'], period)
        # 同一矿的多个通风井只采集一次，报告归属仍由报告内容核对。
        async with self.lock:
            if key not in self.cache:
                self.cache[key] = await self._collect(mine, period, Path(receipt_dir))
            return self.cache[key]

    async def _collect(self, mine, period, folder):
        started = time.monotonic()
        _, mines_body = await self._request('GET', '/mines', params={'with_tree': 0})
        rows = _data(mines_body).get('mines')
        if not isinstance(rows, list) or any(not isinstance(row, dict) for row in rows):
            raise ValueError('资料服务缺少 data.mines 数组')
        matches = [row for row in rows if str(row.get('dept_id')) == str(mine['dept_id'])]
        if len(matches) != 1 or matches[0].get('mine_name') != mine['mine_name']:
            raise ValueError('资料服务与 deptTree 的矿井 ID/全称不一致，停止采集以免取错矿')
        name, kind = _part(matches[0]['mine_name']), _part(self.settings.report_type)
        base = {'mine_name': name, 'dept_id': str(mine['dept_id']), 'year': int(period[:4])}
        status, sync = await self._request('POST', '/doc-types/sync', json=base)
        if status == 409:
            raise ValueError('资料类型同步仍在运行，请稍后恢复本批次')
        _check_success(sync)
        _, types = await self._request('GET', '/doc-types', params={'mine': name, 'year': base['year'], 'name': '通风机'})
        rows = _data(types).get('items')
        if not isinstance(rows, list) or any(not isinstance(row, dict) for row in rows):
            raise ValueError('资料类型接口缺少 data.items 数组')
        matches = [row for row in rows if row.get('level3') == kind]
        if len(matches) != 1 or not str(matches[0].get('file_name_id', '')).isdigit():
            raise ValueError('未唯一找到通风机测定报告资料类型，不能使用其他矿的类型编号')
        request = {**base, 'file_name_ids': [int(matches[0]['file_name_id'])], 'period': period, 'dry_run': False}
        record_path = folder / 'task.json'
        previous = None
        if record_path.is_file():
            previous = json.loads(record_path.read_text(encoding='utf-8'))
            if previous.get('request') != request:
                raise ValueError('本批次已冻结的报告采集参数发生变化，请新建批次')
        task_id = previous.get('task_id') if previous else None
        if not previous:
            _write(record_path, {'request': request, 'task_id': None, 'status': 'pending'})
        deadline = time.monotonic() + self.settings.task_timeout
        busy_logged = False
        final = None
        while True:
            if time.monotonic() >= deadline:
                raise ValueError('报告采集等待超时；采集参数及已取得的任务 ID 已保留，可恢复本批次')
            if not task_id:
                status, body = await self._request('POST', '/crawl', json=request)
                if status == 409:
                    # 文档未保证 409 返回任务 ID；等待后仅重试同一幂等请求。
                    if not busy_logged:
                        logger.info('资料服务忙，等待现有采集任务结束', extra={'kv': {'event': 'report.crawl.busy', 'period': period}})
                        busy_logged = True
                    await asyncio.sleep(self.settings.poll_seconds)
                    continue
                task_id = _data(body).get('task_id')
                if not isinstance(task_id, str) or not re.fullmatch(r'[A-Za-z0-9_-]{1,128}', task_id):
                    raise ValueError('采集接口缺少有效 task_id')
                _write(record_path, {'request': request, 'task_id': task_id, 'status': 'running'})
                logger.info('报告采集已启动', extra={'kv': {'event': 'report.crawl.started',
                    'crawl_task_id': task_id, 'period': period, 'report_type': kind}})
            _, final = await self._request('GET', '/crawl/' + quote(task_id, safe=''))
            state = _data(final).get('status')
            _write(record_path, {'request': request, 'task_id': task_id, 'status': state, 'response': final})
            if state in {'done', 'partial', 'failed'}:
                if state != 'done':
                    raise ValueError(f'报告采集终态为 {state}，详见 report-source/task.json，不能当作完整资料分析')
                break
            if state not in {'queued', 'pending', 'running'}:
                raise ValueError('报告采集返回未知任务状态，详见 report-source/task.json')
            await asyncio.sleep(self.settings.poll_seconds)
        # 快速调用文档未定义返回文件清单，按其约定的唯一矿/年/类型目录读取。
        # 集合可能包含往期原件；记录为年度已保存集合，不冒称本期新增附件。
        records = await asyncio.to_thread(self._files, name, base['year'], kind)
        attach = _data(final).get('attach')
        attachment_count = attach if type(attach) is int else len(attach) if isinstance(attach, list) else None
        if attachment_count and not records:
            raise ValueError('采集任务返回有附件，但挂载目录没有PDF，请检查目录映射及附件格式')
        receipt = {'request': request, 'task_id': task_id, 'response': final,
            'collection_scope': '同矿同年度同资料类型目录的已保存PDF集合，可能包含往期原件；报告日期及井/风机归属须核对原文',
            'files': [{key: val for key, val in row.items() if key != '_local_source'} for row in records]}
        _write(folder / 'files.json', receipt)
        logger.info('报告采集完成', extra={'kv': {'event': 'report.crawl.completed', 'crawl_task_id': task_id,
            'period': period, 'pdf_count': len(records), 'duration_ms': round((time.monotonic() - started) * 1000)}})
        return records, {'source': 'crawler_mount', **request}, receipt

    def _files(self, mine_name, year, kind):
        mount = Path(self.settings.source_dir)
        if not mount.is_dir():
            raise ValueError('报告挂载目录不存在，请挂载 /home/chief/source_data 并配置 MAIN_FAN_REPORT_SOURCE_DIR')
        root = mount.resolve()
        source = root / mine_name / str(year) / kind
        if any(part.is_symlink() for part in (source, source.parent, source.parent.parent)):
            raise ValueError('报告矿井/年份/类型目录不能是符号链接')
        if not source.exists():
            return []
        if not source.is_dir() or not source.resolve().is_relative_to(root):
            raise ValueError('报告资料目录不是挂载范围内的有效目录')
        records = []
        for path in sorted(source.iterdir()):
            if path.suffix.lower() != '.pdf':
                continue
            if path.is_symlink() or not path.is_file():
                raise ValueError('报告目录存在无效 PDF 或符号链接')
            if len(records) >= self.settings.max_files:
                raise ValueError('报告数量超过 MAIN_FAN_REPORT_MAX_FILES，不截断报告集合')
            if not 0 < path.stat().st_size <= self.settings.max_file_mb * 1024 * 1024:
                raise ValueError('报告文件为空或超过配置大小上限')
            records.append({'filename': path.name, 'is_mock': False, 'binding': 'mine_directory_content_unverified',
                'note': '来自核对ID和全称后的矿级资料目录；仍须核对PDF正文的报告日期、通风井及风机归属。',
                'source_relative_path': str(path.relative_to(root)), '_local_source': str(path)})
        return records

    async def download(self, record, target):
        def copy():
            # 文件在等待期间也可能变化；复制时重新检查边界、大小及签名。
            source = Path(record['_local_source'])
            root = Path(self.settings.source_dir).resolve()
            expected = root / record['source_relative_path']
            if (not expected.is_relative_to(root) or source.is_symlink()
                    or source.resolve() != expected or not source.is_file()):
                raise ValueError('报告原件不在有效挂载目录内')
            target.parent.mkdir(parents=True, exist_ok=True)
            partial = target.with_suffix('.part')
            sha, size = hashlib.sha256(), 0
            try:
                with source.open('rb') as incoming, partial.open('wb') as outgoing:
                    before = os.fstat(incoming.fileno())
                    if incoming.read(5) != b'%PDF-':
                        raise ValueError('挂载报告不是 PDF 原件')
                    incoming.seek(0)
                    while chunk := incoming.read(1024 * 1024):
                        size += len(chunk)
                        if size > self.settings.max_file_mb * 1024 * 1024:
                            raise ValueError('复制过程中报告超过配置大小上限')
                        sha.update(chunk)
                        outgoing.write(chunk)
                    after = source.stat()
                    if ((before.st_dev, before.st_ino, before.st_size, before.st_mtime_ns)
                            != (after.st_dev, after.st_ino, after.st_size, after.st_mtime_ns)
                            or size != after.st_size or source.resolve() != expected):
                        raise ValueError('报告复制过程中原件发生变化，请稍后重试')
                partial.replace(target)
            finally:
                partial.unlink(missing_ok=True)
            logger.debug('PDF 已复制到 workspace', extra={'kv': {'event': 'report.copied',
                'filename': source.name, 'size': size, 'sha256': sha.hexdigest(), 'target': str(target)}})
            return sha.hexdigest()
        return await asyncio.to_thread(copy)