Crawl/backend/app/utils/price_parser.py

"""
price_parser.py
统一价格结构，所有模型类型输出相同字段：

  tier_min     : 档位下限（token数 或 秒数），0 表示从0开始，None 表示无档位
  tier_max     : 档位上限（token数 或 秒数），None 表示无上限
  tier_unit    : 档位单位，"tokens" 或 "seconds"，None 表示无档位
  input_price  : 输入价格（元/百万tokens 或 0），视频/图像为 0
  output_price : 输出价格（元/百万tokens）或视频/图像的生成价格
  currency     : "CNY"
  unit         : 计费单位原始字符串
  label        : 原始 key

视频规格 -> 秒数映射：
  480P  -> 0 ~ 480
  720P  -> 0 ~ 720  (或 481 ~ 720)
  1080P -> 0 ~ 1080 (或 721 ~ 1080)
  4K    -> 0 ~ 2160
"""

from __future__ import annotations

import re
from typing import Any, Dict, List, Optional


# ── 视频规格 -> 最大秒数 ────────────────────────────────────────────────────────
_VIDEO_SPEC_MAX: Dict[str, int] = {
    "480P":  480,
    "480p":  480,
    "720P":  720,
    "720p":  720,
    "1080P": 1080,
    "1080p": 1080,
    "2K":    1440,
    "4K":    2160,
}

# 非 token 计费单位
_NON_TOKEN_UNITS = re.compile(r"每秒|每张|每次|每帧|/秒|/张|/次", re.I)

# token 阶梯 key 正则
# 情况1：input<=128k / 32k<input<=128k（有上限）
_TIER_RE = re.compile(
    r"^(?:([\d.]+[KkMm]?)\s*<\s*)?(?:input|输入)\s*<=?\s*([\d.]+[KkMm]?)$",
    re.I,
)
# 情况2：256k<input（只有下限，无上限）
_TIER_RE_LO_ONLY = re.compile(
    r"^([\d.]+[KkMm]?)\s*<\s*(?:input|输入)$",
    re.I,
)


def _to_tokens(val: str) -> Optional[int]:
    """把 '32k'/'128K'/'1M' 转成 token 整数。"""
    s = str(val).strip().upper().replace(",", "")
    m = re.match(r"^([\d.]+)\s*([KMG]?)$", s)
    if not m:
        return None
    num = float(m.group(1))
    suffix = m.group(2)
    if suffix == "K":
        return int(num * 1_000)
    if suffix == "M":
        return int(num * 1_000_000)
    return int(num)


def _parse_price(obj: Any) -> Optional[float]:
    if isinstance(obj, (int, float)):
        return float(obj)
    if not isinstance(obj, dict):
        return None

    # 优先取 price_original（原价），爬虫把折扣价和原价分开存
    for key in ("price_original", "price", "price_current"):
        v = obj.get(key)
        if v is not None:
            try:
                return float(v)
            except (TypeError, ValueError):
                continue

    # 嵌套结构：{"8元": {"price": 8}, "24元": {"price": 24}}，取最大
    candidates: List[float] = []
    for sub_v in obj.values():
        if isinstance(sub_v, dict):
            p = _parse_price(sub_v)
            if p is not None:
                candidates.append(p)
        elif isinstance(sub_v, list):
            for item in sub_v:
                if isinstance(item, dict):
                    p = _parse_price(item)
                    if p is not None:
                        candidates.append(p)
    if candidates:
        return max(candidates)
    return None


def _parse_unit(obj: Any) -> Optional[str]:
    if isinstance(obj, dict):
        # 直接有 unit 字段
        u = obj.get("unit")
        if u:
            return u
        # 新格式：嵌套结构，递归搜索
        for sub_v in obj.values():
            if isinstance(sub_v, dict):
                u = sub_v.get("unit")
                if u:
                    return u
            elif isinstance(sub_v, list):
                for item in sub_v:
                    if isinstance(item, dict):
                        u = item.get("unit")
                        if u:
                            return u
    return None


def _parse_tier_key(key: str):
    """解析 token 阶梯 key，返回 (min_tokens, max_tokens) 或 None。"""
    k = key.strip().lower().replace(" ", "")
    m = _TIER_RE.match(k)
    if m:
        lo_str, hi_str = m.group(1), m.group(2)
        lo = _to_tokens(lo_str) if lo_str else 0
        hi = _to_tokens(hi_str) if hi_str else None
        return (lo, hi)
    # 只有下限：256k<input
    m2 = _TIER_RE_LO_ONLY.match(k)
    if m2:
        lo = _to_tokens(m2.group(1))
        return (lo, None)
    return None


def _extract_video_spec(label: str) -> Optional[str]:
    """从 label 中提取视频规格，如 '视频生成（720P）' -> '720P'。"""
    m = re.search(r"[（(]([^）)]+)[）)]", label)
    if m:
        spec = m.group(1).strip()
        if spec.upper() in {k.upper() for k in _VIDEO_SPEC_MAX}:
            return spec.upper()
    # 直接在 label 里找
    for spec in _VIDEO_SPEC_MAX:
        if spec.upper() in label.upper():
            return spec.upper()
    return None


def _build_video_tiers(items: List[Dict]) -> List[Dict]:
    """
    把多个视频规格条目转成连续区间：
    720P(0.6) + 1080P(1.0) ->
      [0, 720,  input=0, output=0.6]
      [721, 1080, input=0, output=1.0]
    """
    # 按 tier_max 排序
    sorted_items = sorted(items, key=lambda x: x["_spec_max"])
    result = []
    prev_max = 0
    for item in sorted_items:
        spec_max = item["_spec_max"]
        result.append({
            "label": item["label"],
            "tier_min": prev_max + (1 if prev_max > 0 else 0),
            "tier_max": spec_max,
            "tier_unit": "seconds",
            "input_price": 0.0,
            "output_price": item["price"],
            "currency": item["currency"],
            "unit": item["unit"],
        })
        prev_max = spec_max
    return result


def parse_prices(prices: Dict[str, Any]) -> List[Dict]:
    result: List[Dict] = []
    video_items: List[Dict] = []
    input_entry: Optional[Dict] = None

    _CACHE_RE = re.compile(r"缓存|cache", re.I)

    # 预扫描：是否存在阶梯 key
    _has_tiers = any(_parse_tier_key(k) is not None for k in prices)

    # 有阶梯时，预先提取顶层 "output"/"输出" 里的所有独立 output 价格
    # 爬虫把所有 output 价格汇总在顶层 output key 里，需要按大小分配给各阶梯
    _output_prices: List[float] = []
    if _has_tiers:
        for _ok in ("output", "输出"):
            if _ok in prices and isinstance(prices[_ok], dict):
                for _sub_v in prices[_ok].values():
                    _p = _parse_price(_sub_v)
                    if _p is not None:
                        _output_prices.append(_p)
                break
        _output_prices.sort()  # 从小到大，与阶梯从低到高对应

    for key, val in prices.items():
        # 跳过缓存相关条目
        if _CACHE_RE.search(key):
            continue

        if not isinstance(val, dict):
            continue

        # ── token 阶梯 ──
        tier = _parse_tier_key(key)
        if tier is not None:
            entry: Dict = {
                "label": key,
                "tier_min": tier[0],
                "tier_max": tier[1],
                "tier_unit": "tokens",
                "input_price": None,
                "output_price": None,
                "currency": "CNY",
                "unit": None,
            }
            for sub_key, sub_val in val.items():
                sk = sub_key.strip()
                if _CACHE_RE.search(sk):
                    continue
                price = _parse_price(sub_val)
                unit = _parse_unit(sub_val)
                if unit:
                    entry["unit"] = unit
                if re.match(r"^输入|^input", sk, re.I):
                    entry["input_price"] = price
                elif re.match(r"^输出|^output", sk, re.I):
                    entry["output_price"] = price
                elif price is not None and entry["input_price"] is None:
                    # 无标签的纯价格 → 当作输入价格
                    entry["input_price"] = price
            result.append(entry)
            continue

        # 有阶梯时，顶层 output 已预提取，跳过
        if _has_tiers and re.match(r"^(输出|output)$", key.strip(), re.I):
            continue

        price = _parse_price(val)
        unit = _parse_unit(val)

        # ── 视频/图像按单位计费 ──
        if _NON_TOKEN_UNITS.search(unit or ""):
            spec = _extract_video_spec(key)
            if spec and spec in _VIDEO_SPEC_MAX:
                video_items.append({
                    "label": key,
                    "_spec_max": _VIDEO_SPEC_MAX[spec],
                    "price": price,
                    "currency": "CNY",
                    "unit": unit,
                })
            else:
                result.append({
                    "label": key,
                    "tier_min": None,
                    "tier_max": None,
                    "tier_unit": None,
                    "input_price": 0.0,
                    "output_price": price,
                    "currency": "CNY",
                    "unit": unit,
                })
            continue

        # ── 简单非阶梯：输入 → 暂存，输出 → 配对 ──
        if re.match(r"^输入|^input", key.strip(), re.I):
            input_entry = {"price": price, "unit": unit}
            continue
        if re.match(r"^输出|^output", key.strip(), re.I):
            result.append({
                "label": "input/output",
                "tier_min": None,
                "tier_max": None,
                "tier_unit": None,
                "input_price": input_entry["price"] if input_entry else None,
                "output_price": price,
                "currency": "CNY",
                "unit": unit or (input_entry["unit"] if input_entry else None),
            })
            input_entry = None
            continue

        # 其他普通标签
        result.append({
            "label": key,
            "tier_min": None,
            "tier_max": None,
            "tier_unit": None,
            "input_price": price,
            "output_price": None,
            "currency": "CNY",
            "unit": unit,
        })

    # 有阶梯时，把预提取的 output 价格按顺序分配给各阶梯
    if _has_tiers and _output_prices:
        # 按 tier_min 从小到大排序（与 output 价格从小到大对应）
        tier_entries = [r for r in result if r.get("tier_unit") == "tokens"]
        tier_entries.sort(key=lambda e: e["tier_min"] or 0)
        for i, entry in enumerate(tier_entries):
            if i < len(_output_prices) and entry["output_price"] is None:
                entry["output_price"] = _output_prices[i]

    # 只有输入没有输出 → 单独一条
    if input_entry:
        result.append({
            "label": "input",
            "tier_min": None,
            "tier_max": None,
            "tier_unit": None,
            "input_price": input_entry["price"],
            "output_price": None,
            "currency": "CNY",
            "unit": input_entry["unit"],
        })

    # 视频条目转成连续区间
    if video_items:
        result.extend(_build_video_tiers(video_items))

    return result