Auto-Finance/src/pipeline.py

"""
报销全流程编排

将发票提取 → 差旅/普通信息提取 → 浏览器填报串联为一条管道，
数据在内存中流转，同时生成 CSV 中间产物。

发票类型区分：
    - 差旅发票（train/hotel）：不生成易耗品出库单，走差旅报销流程
    - 普通发票（general）：生成易耗品出库单，走普通报销流程

数据流变更（2026-06-11）：
    在发票提取和匹配完成后立即判断报销类型（差旅/普通），
    差旅调用 LLM 提取 travel_info.json，普通预留 normal_info.json。
    Bot 仅负责接收信息并填报，不再承担信息提取职责。
"""

import json
from pathlib import Path
from typing import Any, cast

from . import get_logger
from .config import load_config
from .doc.extractor import extract_invoices
from .doc.invoice import (
    classify_invoice_batch,
    save_application_json,
    save_invoice_csv,
)
from .doc.invoice import (
    save_csv as save_payment_csv,
)
from .doc.llm_extractor import (
    CACHE_DIR_NAME,
    extract_normal_info,
    extract_travel_info,
    load_cache,
)

log = get_logger("pipeline")


def _classify_from_cache(cache_path: Path) -> dict[str, list[dict[str, Any]]]:
    """从缓存目录读取发票数据并按类型分组"""
    from .doc.llm_extractor import load_cache

    cache_map = load_cache(cache_path)
    invoices = [data for data in cache_map.values() if data.get("invoice_type") not in ("application", "payment")]
    return classify_invoice_batch(invoices)


def _extract_travel_info_if_needed(groups: dict[str, list[dict[str, Any]]], cache_path: Path) -> dict[str, Any] | None:
    """当存在差旅发票时，调用 LLM 提取差旅信息并缓存。

    Returns:
        差旅信息字典，非差旅时返回 None。
    """
    if not groups.get("travel"):
        return None

    from .doc.llm_extractor import CACHE_DIR_NAME, load_cache

    # 检查缓存是否已有
    cache_map = load_cache(cache_path)
    if cache_map.get("travel_info"):
        log.info("使用已有差旅信息缓存")
        return cast(dict[str, Any] | None, cache_map["travel_info"])

    log.info("开始提取差旅信息...")
    travel_info = extract_travel_info(source_dir=cache_path)

    # 保存到缓存
    cache_dir = cache_path / CACHE_DIR_NAME
    cache_dir.mkdir(parents=True, exist_ok=True)

    with open(cache_dir / "travel_info.json", "w", encoding="utf-8") as f:
        json.dump(travel_info, f, ensure_ascii=False, indent=2)
    log.info("差旅信息已保存到缓存")
    return travel_info


def _extract_normal_info_if_needed(groups: dict[str, list[dict[str, Any]]], cache_path: Path) -> dict[str, Any] | None:
    """当存在普通发票时，调用 LLM 提取普通报销信息并缓存。

    Returns:
        普通报销信息字典，非普通时返回 None。
    """
    if not groups.get("general"):
        return None

    # 检查缓存是否已有
    cache_map = load_cache(cache_path)
    if cache_map.get("normal_info"):
        log.info("使用已有普通发票信息缓存")
        return cast(dict[str, Any] | None, cache_map["normal_info"])

    log.info("开始提取普通发票信息...")
    normal_info = extract_normal_info(source_dir=cache_path)

    # 保存到缓存
    cache_dir = cache_path / CACHE_DIR_NAME
    cache_dir.mkdir(parents=True, exist_ok=True)

    with open(cache_dir / "normal_info.json", "w", encoding="utf-8") as f:
        json.dump(normal_info, f, ensure_ascii=False, indent=2)
    log.info("普通发票信息已保存到缓存")
    return normal_info


def run_pipeline(
    step: str = "all",
    username: str | None = None,
    password: str | None = None,
    cache_dir: str | None = None,
) -> int:
    """执行报销流程

    Args:
        step: all | invoice | submit
        username: 覆盖 config.json 中的用户名
        password: 覆盖 config.json 中的密码
        cache_dir: 发票缓存目录（包含 .invoice_cache 子目录）
    """
    config = load_config()
    if username:
        config["username"] = username
    if password:
        config["password"] = password

    project_dir = Path(__file__).parent.parent
    cache_path = Path(cache_dir) if cache_dir else project_dir / "scripts" / "data"

    # --------------------------------------------------
    # Step 1: 发票提取 + 类型判断 + 信息提取
    # --------------------------------------------------
    payment_records: list[dict[str, str]] | None = None
    groups: dict[str, list[dict[str, str]]] | None = None
    travel_info: dict[str, Any] | None = None
    normal_info: dict[str, Any] | None = None

    if step in ("all", "invoice"):
        log.info("=" * 60)
        log.info("[1/2] 发票提取")
        log.info("=" * 60)

        payment_records, applications, groups = extract_invoices(str(cache_path))
        if not payment_records:
            log.error("未提取到任何发票数据")
            return 1

        save_payment_csv(payment_records, cache_path / "payment_records.csv")
        save_invoice_csv(payment_records, cache_path / "invoice_summary.csv")

        if applications:
            save_application_json(applications, cache_path / "travel_applications.json")

        log.info(f"发票分类: 差旅 {len(groups['travel'])} 张, 普通 {len(groups['general'])} 张")

        # 发票提取完成后立即判断类型
        is_travel = bool(groups["travel"]) and not bool(groups["general"])
        if is_travel:
            travel_info = _extract_travel_info_if_needed(groups, cache_path)
        else:
            normal_info = _extract_normal_info_if_needed(groups, cache_path)

        if step == "invoice":
            log.info("[1/2] 发票提取 完成")
            return 0

    # --------------------------------------------------
    # Step 2: 浏览器填报
    # --------------------------------------------------
    if step in ("all", "submit"):
        log.info("=" * 60)
        log.info("[2/2] 报销提交")
        log.info("=" * 60)

        from .bot import run_bot

        if groups is None:
            groups = _classify_from_cache(cache_path)

        is_travel = bool(groups["travel"]) and not bool(groups["general"])
        if is_travel:
            if travel_info is None:
                travel_info = _extract_travel_info_if_needed(groups, cache_path)
            log.info("检测到纯差旅发票，使用差旅报销模式")
            run_bot(config, work_dir=cache_path, travel_info=travel_info)
        else:
            if normal_info is None:
                normal_info = _extract_normal_info_if_needed(groups, cache_path)
            log.info("检测到普通发票，使用普通报销模式")
            run_bot(config, work_dir=cache_path, normal_info=normal_info)

        if step == "submit":
            log.info("[2/2] 报销提交 完成")
            return 0

    log.info("=" * 60)
    log.info("全流程执行完毕")
    log.info("=" * 60)
    return 0