#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""资金缺口「讲人话」— G2 验收件（P0-A）

任务：fullstack-20261017-warning-capability-build-p0（上游 yanxue-20261017-warning-capability-build-p0）
方案：`_reports/yanxue-20261017-warning-capability-build-p0.md` §G1/G2 + §四「数字零错误（可复算）」

纪律（本脚本存在的唯一理由）：
  1. **数字由规则层算死**：所有数字一律来自 `forecast_cash_flow_with_plans` 的返回值，
     脚本把这份 payload 作为「白名单」喂给 LLM，并**事后断言**生成文本里的数字集合
     ⊆ 规则层数字集合（数值相等，非字符串相等；`20` 与 `20.0` 视为同一数字）。
  2. **LLM 只写文字**：prompt 明确禁止四舍五入、估算、派生、新增数字、自造日期。
  3. **断言失败即不带文案上线**：首次违例 → 收紧 prompt 重试一次；仍违例 → 退出码 1、
     `assertion.pass=false`、`narrative` 置空（调用方不得把文案推给用户）。

用法：
  # 规则层取数（只读，走 cma_reader）→ LLM 出文案 → 数字守恒断言
  python3 backend/scripts/cash_gap_narrative.py --entity 1 --days 30
  # 只跑断言自测（不连库、不连 LLM，离线可判）
  python3 backend/scripts/cash_gap_narrative.py --selftest
  # 用现成 payload 出文案（离线复算/回放）
  python3 backend/scripts/cash_gap_narrative.py --payload-json p.json
退出码：0 = 断言通过且文案完整；1 = 数字违例 / 段落缺失 / 取数或调用失败
"""
from __future__ import annotations

import argparse
import asyncio
import json
import os
import re
import sys
from datetime import datetime
from decimal import Decimal, InvalidOperation

BACKEND_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
SCRIPTS_DIR = os.path.dirname(os.path.abspath(__file__))
for _p in (BACKEND_DIR, SCRIPTS_DIR):
    if _p not in sys.path:
        sys.path.insert(0, _p)

READER_ENV_DEFAULT = "/root/.hermes/profiles/wecom-fullstack/.secrets/cma-reader.env"

#: DeepSeek Key 来源候选（按序取第一个「非空」；脚本只记来源路径，不打印 Key）
#: 生产 CMA backend/.env 中的 Key 已被供应商判为 invalid（2026-10-17 实测 401），
#: 故 Bot 侧优先用 Hermes 通道 Key；生产 Key 轮换另走派单，本脚本不改生产配置。
KEY_ENV_CANDIDATES = (
    "/root/.hermes/profiles/wecom-fullstack/.env",
    "/root/.hermes/.env",
    "/root/cma-management/backend/.env",
)

#: 文本里允许出现的非数字词表（避免把「TOP3」里的 3 当数字）——prompt 亦禁止
DATE_PATTERNS = (
    re.compile(r"\d{4}-\d{2}-\d{2}"),
    re.compile(r"\d{4}年\d{1,2}月\d{1,2}日"),
    re.compile(r"\d{4}-\d{2}(?!\d)"),
    re.compile(r"\d{4}年\d{1,2}月(?!\d)"),
)
NUMBER_PATTERN = re.compile(r"-?\d+(?:\.\d+)?")

REQUIRED_SECTIONS = ("原因", "建议")


# ────────────────────────── 数字集合工具 ──────────────────────────
def _dec(value) -> Decimal | None:
    """把任意数值转 Decimal（量化到 2 位，与规则层 round(2) 口径一致）；非法返回 None。"""
    try:
        d = Decimal(str(value))
    except (InvalidOperation, ValueError):
        return None
    if not d.is_finite():
        return None
    return d.quantize(Decimal("0.01"))


def collect_rule_numbers(payload) -> tuple[set[Decimal], set[str]]:
    """递归收集 payload 中所有数值与日期（白名单）。

    · 数值同时收 x 与 -x：文本用「下降 37.81%」表达负趋势属同一数字，不算改数
    · 布尔值不算数字；字符串里的日期单独收集，不计入数字白名单
    """
    values: set[Decimal] = set()
    dates: set[str] = set()

    def walk(node):
        if isinstance(node, bool) or node is None:
            return
        if isinstance(node, (int, float, Decimal)):
            d = _dec(node)
            if d is not None:
                values.add(d)
                values.add(-d)
            return
        if isinstance(node, str):
            for pat in DATE_PATTERNS:
                dates.update(pat.findall(node))
            return
        if isinstance(node, dict):
            for v in node.values():
                walk(v)
            return
        if isinstance(node, (list, tuple)):
            for v in node:
                walk(v)

    walk(payload)
    return values, dates


def extract_text_numbers(text: str) -> tuple[list[tuple[str, Decimal]], set[str], str]:
    """抽取文本中的日期与数字；返回 (数字token列表, 日期集合, 去日期后的文本)。"""
    dates: set[str] = set()
    stripped = text
    for pat in DATE_PATTERNS:
        for m in pat.findall(stripped):
            dates.add(m)
        stripped = pat.sub(" ", stripped)

    numbers: list[tuple[str, Decimal]] = []
    for raw in NUMBER_PATTERN.findall(stripped):
        d = _dec(raw)
        if d is not None:
            numbers.append((raw, d))
    return numbers, dates, stripped


def assert_number_conservation(text: str, payload) -> dict:
    """核心断言：文本数字 ⊆ 规则层数字；日期 ⊆ 规则层日期。返回结构化结果。"""
    allowed_values, allowed_dates = collect_rule_numbers(payload)
    numbers, dates, _ = extract_text_numbers(text)

    value_violations = [
        {"token": tok, "value": str(val)}
        for tok, val in numbers
        if val not in allowed_values
    ]
    date_violations = sorted(d for d in dates if d not in allowed_dates)

    return {
        "pass": not value_violations and not date_violations,
        "checked_numbers": len(numbers),
        "allowed_numbers": len(allowed_values),
        "value_violations": value_violations,
        "date_violations": date_violations,
        "allowed_values_sample": sorted(str(v) for v in allowed_values)[:40],
        "allowed_dates": sorted(allowed_dates),
    }


# ────────────────────────── 规则层取数 ──────────────────────────
def _dotenv(path: str) -> dict:
    env = {}
    for line in open(path, encoding="utf-8"):
        line = line.strip()
        if line and "=" in line and not line.startswith("#"):
            k, v = line.split("=", 1)
            env[k.strip()] = v.strip().strip('"').strip("'")
    return env


def resolve_deepseek_key() -> tuple[str, str]:
    """按候选顺序取 DeepSeek Key；返回 (key, 来源路径)。只用于调用，不落盘不打印。"""
    current = os.environ.get("DEEPSEEK_API_KEY", "").strip()
    for path in KEY_ENV_CANDIDATES:
        if not os.path.exists(path):
            continue
        try:
            key = _dotenv(path).get("DEEPSEEK_API_KEY", "").strip()
        except OSError:
            continue
        if key:
            return key, path
    return current, "environment" if current else ""


def load_rule_payload(entity_id: int, days: int) -> dict:
    """只读取数：cma_reader 账号（SELECT only）→ 引擎同一函数，与端点同源。"""
    from sqlalchemy import create_engine
    from sqlalchemy.orm import sessionmaker

    from app.utils import cash_forecast_engine as engine

    rd = _dotenv(os.environ.get("CMA_READER_ENV", READER_ENV_DEFAULT))
    url = ("mysql+pymysql://%s:%s@%s:%s/%s?charset=utf8mb4"
           % (rd["CMA_READER_DB_USER"], rd["CMA_READER_DB_PASS"],
              rd["CMA_READER_DB_HOST"], rd["CMA_READER_DB_PORT"], rd["CMA_READER_DB_NAME"]))
    eng = create_engine(url, pool_pre_ping=True)
    db = sessionmaker(bind=eng)()
    try:
        payload = engine.forecast_cash_flow_with_plans(entity_id, db, days=days)
    finally:
        db.close()
    return payload


# ────────────────────────── prompt ──────────────────────────
SYSTEM_PROMPT = (
    "你是管理会计，只负责把已经算好的资金缺口结论写成老板能读懂的「原因」和「建议」。"
    "硬性规则："
    "① 只能引用我给你的数字，逐字照抄，不得四舍五入、不得换算单位、不得估算、不得派生新数字；"
    "② 禁止出现我未提供的任何新数字（包括数量、比例、天数、排名，如「TOP3」里的 3 不能用）；"
    "③ 日期只能照抄我给的日期；"
    "④ 输出两段，段首分别为「原因：」与「建议：」，每段不超过 120 字；"
    "⑤ 不要复述表格，不要下结论性预测以外的推断。"
)


def build_prompt(payload: dict) -> str:
    """把规则层 payload 压成「可引用数字白名单」，避免把整包原始明细灌给模型。"""
    def r(v):
        return None if v is None else round(float(v), 2)

    first_gap = (payload.get("gap_dates") or [None])[0] or {}
    first_level = {"red": "红灯", "yellow": "黄灯"}.get(str(first_gap.get("level") or ""), "无")

    compact = {
        "账套entity_id": payload.get("entity_id"),
        "预测窗口天数": payload.get("days"),
        "当前现金余额(万元)": r(payload.get("base_cash")),
        "黄灯警戒线(万元)": r(payload.get("warning_line")),
        "红灯警戒线(万元)": r(payload.get("critical_line")),
        "阈值来源": (payload.get("threshold_source") or {}).get("source"),
        "最低余额(万元)": r(payload.get("min_cash")),
        "最低余额日期": payload.get("min_cash_date"),
        "缺口天数": len(payload.get("gap_dates") or []),
        "首个缺口": first_gap or None,
        "首个缺口级别": first_level,
        "月度经营现金流预算(万元/月)": r((payload.get("trends") or {}).get("budget_monthly_ocf")),
        "经营现金流趋势(%)": r((payload.get("trends") or {}).get("operating_cash_flow_trend_pct")),
        "应收趋势(%)": r((payload.get("trends") or {}).get("receivables_trend_pct")),
        "应付趋势(%)": r((payload.get("trends") or {}).get("payables_trend_pct")),
        "逾期应收笔数": (payload.get("summary") or {}).get("overdue_receive_count"),
        "逾期应收金额(万元)": r((payload.get("summary") or {}).get("overdue_receive_amount")),
    }
    return (
        "以下是规则层已经算好的资金缺口数据（唯一可引用数字来源）：\n"
        + json.dumps(compact, ensure_ascii=False, indent=2)
        + "\n请只基于这些数字写「原因」与「建议」两段。"
    )


def strict_retry_prompt(payload: dict, bad_text: str, violations: list) -> str:
    return (
        build_prompt(payload)
        + "\n\n上一次输出违反了数字纪律，出现了未授权的数字："
        + json.dumps(violations, ensure_ascii=False)
        + "\n上一次输出为：\n" + bad_text
        + "\n请重写，删除上述未授权数字，其余规则不变。"
    )


# ────────────────────────── 主流程 ──────────────────────────
def split_sections(text: str) -> dict:
    """把模型输出切成 原因 / 建议 两段（段首关键字）。"""
    out = {"reason": "", "advice": ""}
    reason = re.search(r"原因[：:]\s*(.*?)(?=建议[：:]|$)", text, re.S)
    advice = re.search(r"建议[：:]\s*(.*)$", text, re.S)
    if reason:
        out["reason"] = reason.group(1).strip()
    if advice:
        out["advice"] = advice.group(1).strip()
    return out


async def generate(entity_id: int, days: int, payload: dict | None = None) -> dict:
    key, key_source = resolve_deepseek_key()
    if key:
        # 必须在 import ai_brief 之前落位：ai_brief 模块级 load_dotenv 不覆盖已存在的环境变量，
        # 否则会读到生产 .env 里已被判 invalid 的旧 Key（2026-10-17 实测 401）。
        os.environ["DEEPSEEK_API_KEY"] = key

    from ai_brief import call_deepseek  # 复用既有 DeepSeek 通道（scripts/ai_brief.py:24）

    payload = payload if payload is not None else load_rule_payload(entity_id, days)
    prompt = build_prompt(payload)
    text = await call_deepseek(prompt, SYSTEM_PROMPT)
    assertion = assert_number_conservation(text or "", payload)
    attempts = 1

    if not assertion["pass"] and text:
        text2 = await call_deepseek(
            strict_retry_prompt(payload, text, assertion["value_violations"] + assertion["date_violations"]),
            SYSTEM_PROMPT,
        )
        assertion2 = assert_number_conservation(text2 or "", payload)
        attempts = 2
        if assertion2["pass"]:
            text, assertion = text2, assertion2
        else:
            assertion = assertion2  # 以最后一次为准（仍为 fail）

    sections = split_sections(text or "")
    missing = [s for s in REQUIRED_SECTIONS if not sections[{"原因": "reason", "建议": "advice"}[s]]]
    ok = assertion["pass"] and not missing

    return {
        "task": "fullstack-20261017-warning-capability-build-p0",
        "stage": "G2 cash gap narrative (number-conservation asserted)",
        "generated_at": datetime.now().isoformat(timespec="seconds"),
        "entity_id": entity_id,
        "days": days,
        "model": "deepseek-chat",
        "llm_key_source": key_source,
        "llm_empty_response": not bool((text or "").strip()),
        "attempts": attempts,
        "assertion": assertion,
        "missing_sections": missing,
        "pass": ok,
        "narrative": ({"reason": sections["reason"], "advice": sections["advice"], "raw": text}
                      if ok else None),
        "rule_layer": {
            "base_cash": payload.get("base_cash"),
            "warning_line": payload.get("warning_line"),
            "critical_line": payload.get("critical_line"),
            "min_cash": payload.get("min_cash"),
            "min_cash_date": payload.get("min_cash_date"),
            "threshold_source": payload.get("threshold_source"),
            "gap_days": len(payload.get("gap_dates") or []),
        },
    }


# ────────────────────────── 断言自测（离线） ──────────────────────────
SELFTEST_PAYLOAD = {
    "entity_id": 1, "days": 30, "base_cash": -9.3, "warning_line": 20.0, "critical_line": 10.0,
    "min_cash": -8.74, "min_cash_date": "2026-10-18", "threshold_source": {"source": "default"},
    "gap_dates": [{"date": "2026-10-18", "gap_amount": 28.74, "level": "red"}],
    "trends": {"budget_monthly_ocf": 16.67, "receivables_trend_pct": -37.81},
    "summary": {"overdue_receive_count": 18, "overdue_receive_amount": 71.68},
}


def selftest() -> int:
    cases = [
        ("good_basic",
         "原因：现金余额-9.3万元，2026-10-18 降至最低-8.74万元，低于警戒线20.0万元，缺口28.74万元；"
         "逾期应收18笔共71.68万元。建议：优先催收逾期应收，暂停非必要支出。",
         True),
        ("good_sign_abs",
         "原因：应收趋势下降37.81%，月度经营现金流预算16.67万元/月，缺口1天。"
         "建议：跟进2026-10-18 起的资金缺口。",
         True),
        ("bad_rounding",
         "原因：现金余额约-9万元，最低-8.7万元。建议：催收应收。",
         False),
        ("bad_derived",
         "原因：现金余额-9.3万元，缺口28.74万元，约为逾期应收的40%。建议：催收。",
         False),
        ("bad_new_count",
         "原因：现金余额-9.3万元。建议：优先催收TOP3客户。",
         False),
        ("bad_fake_date",
         "原因：现金余额-9.3万元，预计2026-11-01 断流。建议：催收。",
         False),
    ]
    failures = 0
    print("== G2 数字守恒断言自测 ==")
    for name, text, expect in cases:
        res = assert_number_conservation(text, SELFTEST_PAYLOAD)
        verdict = "PASS" if res["pass"] == expect else "FAIL"
        if verdict == "FAIL":
            failures += 1
        print("  %-5s %-16s expect_pass=%-5s got_pass=%-5s  value_violations=%s date_violations=%s"
              % (verdict, name, expect, res["pass"],
                 [v["token"] for v in res["value_violations"]], res["date_violations"]))
    total = len(cases)
    print("RESULT: %s  %d/%d" % ("PASS" if failures == 0 else "FAIL", total - failures, total))
    return 0 if failures == 0 else 1


def main() -> int:
    ap = argparse.ArgumentParser(description="资金缺口讲人话（数字守恒断言）")
    ap.add_argument("--entity", type=int, default=1)
    ap.add_argument("--days", type=int, default=30)
    ap.add_argument("--payload-json", help="已有规则层 payload（跳过取数）")
    ap.add_argument("--out", help="输出 JSON 路径")
    ap.add_argument("--selftest", action="store_true", help="只跑断言自测（离线）")
    args = ap.parse_args()

    if args.selftest:
        return selftest()

    payload = None
    if args.payload_json:
        payload = json.load(open(args.payload_json, encoding="utf-8"))

    result = asyncio.run(generate(args.entity, args.days, payload))
    out = args.out or os.path.join(
        "/root/.hermes/profiles/wecom-fullstack/pending-tasks/_evidence",
        "fullstack-20261017-warning-capability-build-p0",
        "g2-narrative-%s-e%d.json" % (datetime.now().strftime("%Y%m%d-%H%M%S"), args.entity))
    os.makedirs(os.path.dirname(out), exist_ok=True)
    with open(out, "w", encoding="utf-8") as fh:
        json.dump(result, fh, ensure_ascii=False, indent=2)

    print(json.dumps({k: result[k] for k in
                      ("pass", "attempts", "assertion", "missing_sections", "rule_layer")},
                     ensure_ascii=False, indent=2))
    if result["narrative"]:
        print("原因：" + result["narrative"]["reason"])
        print("建议：" + result["narrative"]["advice"])
    print("OUT:", out)
    return 0 if result["pass"] else 1


if __name__ == "__main__":
    raise SystemExit(main())