baodan/api/insurance/poster/tasks.py

"""海报后台任务管理。

使用与 PPT parse_worker 相同的后台线程 + Redis 锁模式。
任务状态通过数据库字段追踪，前端轮询获取进度。
"""
import json
import logging
import threading
import os
from datetime import datetime

from insurance.db.compat import db

logger = logging.getLogger(__name__)

_local_locks: set = set()
_local_locks_guard = threading.Lock()
_redis_locks: set = set()

# 超过此时间仍为 queued/generating 的任务视为过期（秒）
STALE_TASK_TIMEOUT = 600  # 10 分钟


def recover_stale_tasks():
    """启动时恢复过期任务：将长时间 queued/generating 的任务标记为 failed。

    应在应用启动时调用。
    """
    from insurance.models.poster_record import PosterRecord
    from insurance.models.poster_case_upload import PosterCaseUpload

    cutoff = datetime.now().timestamp() - STALE_TASK_TIMEOUT

    # 恢复过期的海报生成任务
    stale_records = PosterRecord.query.filter(
        PosterRecord.task_status.in_(["queued", "generating"]),
        PosterRecord.created_at < datetime.fromtimestamp(cutoff),
    ).all()
    for record in stale_records:
        record.task_status = "failed"
        record.task_error = "任务因服务重启而中断，请重新生成"
        record.finished_at = datetime.now()
        logger.warning(f"恢复过期海报任务: record_id={record.id}")

    # 恢复过期的计划书解析任务
    stale_cases = PosterCaseUpload.query.filter(
        PosterCaseUpload.parse_status.in_(["queued", "parsing"]),
        PosterCaseUpload.created_at < datetime.fromtimestamp(cutoff),
    ).all()
    for case in stale_cases:
        case.parse_status = "failed"
        logger.warning(f"恢复过期解析任务: case_id={case.id}")

    if stale_records or stale_cases:
        db.session.commit()
        logger.info(f"已恢复 {len(stale_records)} 个海报任务, {len(stale_cases)} 个解析任务")


# ─── 计划书解析任务 ───────────────────────────────────────

def start_case_parse_task(app, case_upload_id: int) -> bool:
    """启动后台计划书解析任务，返回是否新启动。"""
    lock_key = f"poster_case_parse:{case_upload_id}"
    if not _acquire_lock(lock_key):
        return False

    thread = threading.Thread(
        target=_run_case_parse_task,
        args=(app, case_upload_id, lock_key),
        daemon=True,
    )
    thread.start()
    return True


def _run_case_parse_task(app, case_upload_id: int, lock_key: str):
    with app.app_context():
        try:
            _execute_case_parse(case_upload_id)
        except Exception as exc:
            logger.error(f"计划书解析任务失败 [{case_upload_id}]: {exc}", exc_info=True)
            _mark_case_failed(case_upload_id, str(exc))
        finally:
            _release_lock(lock_key)
            db.session.remove()


def _execute_case_parse(case_upload_id: int):
    from insurance.models.poster_case_upload import PosterCaseUpload

    record = PosterCaseUpload.query.get(case_upload_id)
    if not record:
        return

    filepath = record.source_file_url
    if not filepath or not os.path.exists(filepath):
        record.parse_status = "failed"
        db.session.commit()
        return

    record.parse_status = "parsing"
    db.session.commit()

    from insurance.ppt.extraction import ExtractionOrchestrator
    orchestrator = ExtractionOrchestrator(use_cache=False)

    import asyncio
    parsed = None

    product_context = {}
    if record.product_snapshot_json:
        try:
            product_context = json.loads(record.product_snapshot_json) or {}
        except (json.JSONDecodeError, TypeError):
            logger.warning("计划书产品快照无效: case_id=%s", case_upload_id)
    plan_type = str(product_context.get("planType") or "savings").lower()
    product_data = product_context.get("productData") or {}
    product_name = (
        product_context.get("productName")
        or product_data.get("displayName")
        or ""
    )
    product_aliases = product_data.get("aliases") or []
    company_id = (
        product_context.get("companyId")
        or product_data.get("companyId")
        or ""
    )

    # 优先使用完整解析（获取利益演示表）
    try:
        result = asyncio.run(orchestrator.extract_plan(
            filepath,
            plan_type=plan_type,
            company_id=company_id,
            product_id=record.product_id or "",
            product_name_hint=product_name,
            product_aliases=product_aliases,
        ))
        if result.status != "error" and result.data:
            data = result.data
            parsed = _map_extract_plan_fields(data, result.plan_type, result.status)
            logger.info(f"使用 extract_plan 解析成功: case_id={case_upload_id}, type={result.plan_type}")
    except Exception as exc:
        logger.warning(f"extract_plan 失败，降级到 extract_for_poster: {exc}")

    # 降级：使用轻量解析
    if not parsed:
        try:
            parsed = asyncio.run(orchestrator.extract_for_poster(filepath))
            logger.info(f"使用 extract_for_poster 降级解析: case_id={case_upload_id}")
        except Exception as exc:
            logger.error(f"extract_for_poster 也失败: {exc}")
            record = PosterCaseUpload.query.get(case_upload_id)
            if record:
                record.parse_status = "failed"
                db.session.commit()
            return

    record = PosterCaseUpload.query.get(case_upload_id)
    if not record:
        return
    parse_status = (parsed.get("meta") or {}).get("status", "failed")
    record.parsed_data = json.dumps(parsed, ensure_ascii=False)
    record.parse_status = parse_status
    db.session.commit()


def _map_extract_plan_fields(data: dict, plan_type: str, status: str) -> dict:
    """将 extract_plan 的完整数据映射为海报前端所需字段结构。

    兼容多种可能的字段命名（LLM 输出不固定）。
    """
    insured = data.get("insured") or {}
    policy = data.get("policy") or {}

    def first_value(*values):
        for value in values:
            if value is not None and value != "":
                return value
        return None

    raw_gender = first_value(insured.get("gender"), data.get("gender"))
    gender_map = {
        "male": "男", "m": "男", "男": "男",
        "female": "女", "f": "女", "女": "女",
    }
    gender = gender_map.get(str(raw_gender).strip().lower(), raw_gender) if raw_gender else None

    raw_currency = first_value(policy.get("currency"), data.get("currency"))
    from insurance.ppt.normalizer import _normalize_currency
    currency = _normalize_currency(raw_currency)

    # 保额：优先读取正式提取结构中的 policy
    sum_assured = (
        policy.get("sum_insured")
        or policy.get("basic_sum_insured")
        or policy.get("sum_assured")
        or data.get("sum_insured")
        or data.get("basic_sum_insured")
        or data.get("sum_assured")
        or data.get("face_amount")
        or data.get("coverage_amount")
    )
    # 缴费年期
    premium_term = (
        policy.get("premium_payment_period")
        or policy.get("premium_term")
        or data.get("premium_term")
        or data.get("payment_period")
        or data.get("paying_period")
        or data.get("premium_payment_term")
    )
    # 年缴保费
    annual_premium = (
        policy.get("annual_premium")
        or policy.get("target_premium")
        or data.get("annual_premium")
        or data.get("premium_amount")
        or data.get("yearly_premium")
    )
    # 保障期限
    coverage_period = (
        policy.get("coverage_period")
        or data.get("coverage_period")
        or data.get("policy_term")
        or data.get("coverage_term")
    )

    # 利益演示表
    benefit_table = (
        data.get("benefit_illustration")
        or data.get("benefit_table")
        or data.get("cash_value_projection")
        or []
    )

    result = {
        "age": first_value(insured.get("age"), data.get("age")),
        "gender": gender,
        "currency": currency,
        "sum_assured": sum_assured,
        "premium_term": premium_term,
        "annual_premium": annual_premium,
        "coverage_period": coverage_period,
        "key_benefits": data.get("key_benefits") or [],
        "benefit_table": benefit_table,
        "plan_type": plan_type,
        "extraction_status": status,
    }
    business_fields = (
        "age", "gender", "currency", "sum_assured",
        "premium_term", "annual_premium", "coverage_period",
    )
    valid_count = sum(
        1 for key in business_fields
        if result.get(key) is not None and result.get(key) != ""
    )
    required_by_type = {
        "savings": ("age", "currency", "annual_premium", "premium_term"),
        "ci": ("age", "currency", "sum_assured"),
        "iul": ("age", "currency", "sum_assured", "annual_premium"),
    }
    required = required_by_type.get(str(plan_type or "").lower(), ("age", "currency"))
    missing = [
        key for key in required
        if result.get(key) is None or result.get(key) == ""
    ]
    mapped_status = "failed" if valid_count == 0 else ("partial" if missing else "parsed")
    result["meta"] = {
        "planType": plan_type,
        "status": mapped_status,
        "sourceStatus": status,
        "missingFields": missing,
        "validFieldCount": valid_count,
        "method": (
            (data.get("_meta") or {}).get("method")
            or (data.get("extraction_meta") or {}).get("method")
            or status
        ),
        "provenance": data.get("_provenance") or data.get("provenance") or {},
        "lowQualityPages": (
            (data.get("_meta") or {}).get("low_quality_pages")
            or (data.get("extraction_meta") or {}).get("lowQualityPages")
            or []
        ),
    }
    return result


def _mark_case_failed(case_upload_id: int, error: str):
    from insurance.models.poster_case_upload import PosterCaseUpload

    record = PosterCaseUpload.query.get(case_upload_id)
    if not record:
        return
    record.parse_status = "failed"
    db.session.commit()


# ─── 锁管理 ────────────────────────────────────────────────

def _acquire_lock(key: str) -> bool:
    """获取任务锁（优先 Redis，降级本地内存）。"""
    redis_key = f"poster_task_lock:{key}"
    try:
        from insurance.db.compat import redis_client
        if redis_client and redis_client.set(redis_key, "1", nx=True, ex=1800):
            _redis_locks.add(key)
            return True
        if redis_client:
            return False
    except Exception:
        pass

    with _local_locks_guard:
        if key in _local_locks:
            return False
        _local_locks.add(key)
        return True


def _release_lock(key: str):
    """释放任务锁。"""
    if key in _redis_locks:
        try:
            from insurance.db.compat import redis_client
            if redis_client:
                redis_client.delete(f"poster_task_lock:{key}")
        except Exception:
            pass
        _redis_locks.discard(key)

    with _local_locks_guard:
        _local_locks.discard(key)