"""自由创作视频 token 计费(移植自 jimeng-clone backend/utils/billing.py)。

Token 预估公式(火山官方):(输入视频时长 + 输出时长) × 宽 × 高 × 帧率 / 1024
单价:元/百万tokens,存 ModelConfig.metadata["pricing"],按 分辨率档 × 是否含视频参考 取价。

⚠️ 预估仅用于前端展示与额度预留;真实费用以火山返回 usage.total_tokens 结算。
预留额外加 RESERVE_BUFFER(ledger 禁超预留扣费,真实 tokens 可能高于预估),
结算时 clamp 到预留额并对差额告警。
"""
from decimal import Decimal, ROUND_HALF_UP

# 分辨率 → 像素映射(火山 Seedance 2.0 API 文档)
RESOLUTION_MAP = {
    # 720p
    ("720p", "16:9"): (1280, 720),
    ("720p", "9:16"): (720, 1280),
    ("720p", "4:3"): (1112, 834),
    ("720p", "1:1"): (960, 960),
    ("720p", "3:4"): (834, 1112),
    ("720p", "21:9"): (1470, 630),
    # 480p
    ("480p", "16:9"): (864, 496),
    ("480p", "9:16"): (496, 864),
    ("480p", "4:3"): (752, 560),
    ("480p", "1:1"): (640, 640),
    ("480p", "3:4"): (560, 752),
    ("480p", "21:9"): (992, 432),
    # 1080p(仅标准档)
    ("1080p", "16:9"): (1920, 1080),
    ("1080p", "9:16"): (1080, 1920),
    ("1080p", "4:3"): (1664, 1248),
    ("1080p", "1:1"): (1440, 1440),
    ("1080p", "3:4"): (1248, 1664),
    ("1080p", "21:9"): (2206, 946),
    # 4k(仅标准档)
    ("4k", "16:9"): (3840, 2160),
    ("4k", "9:16"): (2160, 3840),
    ("4k", "4:3"): (3326, 2494),
    ("4k", "1:1"): (2880, 2880),
    ("4k", "3:4"): (2494, 3326),
    ("4k", "21:9"): (4398, 1886),
}

DEFAULT_FPS = 24

# 预留 = 预估费用 × buffer。ledger 的 charge_reserved_credit 在 actual > reserved 时抛错,
# 而火山真实 tokens 有最低用量限制/输入视频真实时长偏差,可能略高于预估。
RESERVE_BUFFER = Decimal("1.10")


def get_resolution(aspect_ratio: str, tier: str) -> tuple:
    """(tier, aspect_ratio) → (width, height)。非法组合 KeyError fail loud,不静默降级。"""
    key = (tier, aspect_ratio)
    if key not in RESOLUTION_MAP:
        raise KeyError(
            f"不支持的分辨率组合: tier={tier!r}, aspect_ratio={aspect_ratio!r}. "
            f"仅支持 480p/720p/1080p/4k × 16:9/9:16/4:3/1:1/3:4/21:9"
        )
    return RESOLUTION_MAP[key]


def estimate_tokens(
    width: int,
    height: int,
    duration: int,
    fps: int = DEFAULT_FPS,
    input_video_duration: float = 0,
) -> int:
    total_duration = duration + (input_video_duration or 0)
    return round(width * height * fps * total_duration / 1024)


def has_video_reference(references: list) -> bool:
    return any((ref or {}).get("type") == "video" for ref in references or [])


def sum_video_duration(references: list) -> float:
    """输入参考视频总时长(秒),计入 token 公式的输入时长项。"""
    return sum(
        float(ref.get("duration") or 0)
        for ref in references or []
        if (ref or {}).get("type") == "video"
    )


def get_token_price(model_config, with_video_ref: bool, resolution: str) -> Decimal:
    """从 ModelConfig.metadata["pricing"] 取单价(元/百万tokens)。

    先按 resolution 精确键,无则回落 "default";缺 pricing/缺键抛 ValueError fail loud
    (1080p/4k 只有标准档配了键,fast/mini 在提交校验就被拒,不允许静默按 720p 计价)。
    """
    pricing = (model_config.metadata or {}).get("pricing") or {}
    # 1080p/4k 有独立价:缺该档键 = 模型不支持该分辨率(fast/mini),fail loud,
    # 绝不按 default(480p/720p)价静默计费——那是欺骗用户(jimeng _get_token_price 同原则)。
    if resolution in ("1080p", "4k") and resolution not in pricing:
        raise ValueError(f"模型 {model_config.name} 不支持 {resolution}——提交校验应已拦截,不应进到计价")
    tier = pricing.get(resolution) or pricing.get("default")
    if not tier:
        raise ValueError(f"模型 {model_config.name} 未配置 pricing(resolution={resolution})")
    key = "with_ref_video" if with_video_ref else "no_ref_video"
    price = tier.get(key)
    if price is None:
        raise ValueError(f"模型 {model_config.name} pricing 缺 {resolution}/{key} 档单价")
    return Decimal(str(price))


def calculate_cost(tokens: int, price: Decimal) -> Decimal:
    """tokens × 单价(元/百万tokens),保留 2 位小数。"""
    cost = Decimal(str(tokens)) * Decimal(str(price)) / Decimal("1000000")
    return cost.quantize(Decimal("0.01"), rounding=ROUND_HALF_UP)


def estimate_video_cost(
    model_config,
    *,
    aspect_ratio: str,
    resolution: str,
    duration: int,
    references: list,
) -> tuple[int, Decimal]:
    """返回 (预估 tokens, 预估费用元)。供提交预留与前端预估口径对齐。"""
    width, height = get_resolution(aspect_ratio, resolution)
    tokens = estimate_tokens(
        width,
        height,
        duration,
        input_video_duration=sum_video_duration(references),
    )
    price = get_token_price(model_config, has_video_reference(references), resolution)
    return tokens, calculate_cost(tokens, price)


def tokens_to_cost(model_config, tokens: int, *, with_video_ref: bool, resolution: str) -> Decimal:
    """按真实 usage.total_tokens 计价(结算口径,与预估同一张价表)。"""
    price = get_token_price(model_config, with_video_ref, resolution)
    return calculate_cost(tokens, price)