import math from sqlalchemy import select, func from sqlalchemy.ext.asyncio import AsyncSession from app.models.user import User from app.models.credit_record import CreditRecord from app.models.system_config import SystemConfig from app.models.video_engine import VideoEngine from app.models.image_engine import ImageEngine from app.models.credit_ratio import CreditRatio from app.utils.id_gen import generate_id from app.utils.exceptions import InsufficientCreditsError from app.services.credit_record_meta_service import CreditRecordMeta, with_user_snapshot async def calc_text_credits(db: AsyncSession, input_tokens: int, output_tokens: int) -> float: """Calculate text credits based on actual token usage and configurable rate.""" result = await db.execute( select(SystemConfig).where(SystemConfig.key == "text_credits_per_1000_tokens").limit(1) ) config = result.scalar_one_or_none() rate = float(config.value) if config else 1.0 total_tokens = input_tokens + output_tokens return round(total_tokens * rate / 1000, 2) async def _get_credit_ratio( db: AsyncSession, *, gen_type: str, resolution: str, engine_id: str | None = None, ) -> CreditRatio | None: """按引擎精确规则优先获取积分规则;找不到时回退到同类型同参数最高规则。""" gen_type = (gen_type or "").lower().strip() resolution = (resolution or "").strip() engine_id = (engine_id or "").strip() or None if engine_id: result = await db.execute( select(CreditRatio) .where(CreditRatio.gen_type == gen_type) .where(CreditRatio.model_config_id == engine_id) .where(CreditRatio.resolution == resolution) .order_by(CreditRatio.base_credits.desc(), CreditRatio.per_second_credits.desc()) .limit(1) ) ratio = result.scalar_one_or_none() if ratio: return ratio result = await db.execute( select(CreditRatio) .where(CreditRatio.gen_type == gen_type) .where(CreditRatio.resolution == resolution) .order_by(CreditRatio.base_credits.desc(), CreditRatio.per_second_credits.desc()) .limit(1) ) return result.scalar_one_or_none() async def calc_video_credits( db: AsyncSession, duration: int, resolution: str, engine_id: str | None = None, ) -> float: """Calculate video credits using CreditRatio table, with fallback to hardcoded. 查询优先级: 1. gen_type=video + engine_id + resolution 精确规则; 2. gen_type=video + resolution 下 base_credits/per_second_credits 最高规则; 3. 原硬编码默认算法。 """ # 如果engine_id为空,默认查询权重最高的视频引擎积分规则 if not engine_id: video_engines_result = await db.execute( select(VideoEngine.id) .where(VideoEngine.is_active == True) .order_by(VideoEngine.priority.desc()) .limit(1) ) engine_id = video_engines_result.scalar_one_or_none() ratio = await _get_credit_ratio( db, gen_type="video", resolution=resolution, engine_id=engine_id, ) if ratio: return round((ratio.base_credits + ratio.per_second_credits * duration) * ratio.ratio, 2) # Fallback base = 60.0 duration_cost = duration * 2.0 multiplier = {"480p": 1, "1080p": 2, "720p": 1.5}.get(resolution, 1.0) return round((base + duration_cost) * multiplier, 2) def calc_credits(duration: int, resolution: str) -> float: """Legacy: hardcoded credit calculation. Prefer calc_video_credits for new code.""" base = 60.0 duration_cost = duration * 2.0 multiplier = {"480p": 2.5, "1080p": 1.5, "720p": 1.0}.get(resolution, 1.0) return round((base + duration_cost) * multiplier, 2) async def calc_image_credits( db: AsyncSession, image_size: str, engine_id: str | None = None, ) -> float: """Calculate image credits using CreditRatio table, with fallback to hardcoded. 查询优先级: 1. gen_type=image + engine_id + image_size 精确规则; 2. gen_type=image + image_size 下 base_credits/per_second_credits 最高规则; 3. 原硬编码默认算法。 """ # 如果engine_id为空,默认查询权重最高的图片引擎积分规则 if not engine_id: image_engines_result = await db.execute( select(ImageEngine.id) .where(ImageEngine.is_active == True) .order_by(ImageEngine.priority.desc()) .limit(1) ) engine_id = image_engines_result.scalar_one_or_none() ratio = await _get_credit_ratio( db, gen_type="image", resolution=image_size, engine_id=engine_id, ) if ratio: return round(ratio.base_credits * ratio.ratio, 2) # Fallback multiplier = {"4K": 2.0, "2K": 1.0}.get(image_size, 1.0) base_cost = 4.0 return round(base_cost * multiplier, 2) async def _get_existing_credit_record_by_biz_key( db: AsyncSession, *, user_id: str, biz_key: str | None, ) -> CreditRecord | None: """按正式业务幂等键查找已有积分流水。""" if not biz_key: return None result = await db.execute( select(CreditRecord) .where(CreditRecord.user_id == user_id, CreditRecord.biz_key == biz_key) .limit(1) ) return result.scalar_one_or_none() async def deduct_credits( db: AsyncSession, user_id: str, amount: float, description: str, related_id: str | None = None, *, biz_key: str | None = None, refund_for_biz_key: str | None = None, record_meta: CreditRecordMeta | dict | None = None, ) -> User: """扣减用户积分,并写入消费流水。 并发安全点: - 先用 SELECT ... FOR UPDATE 锁住 users 行,避免余额覆盖。 - biz_key 不为空时,作为正式业务幂等键;重复调用直接返回当前用户,不重复扣。 """ amount = round(float(amount or 0), 2) if amount <= 0: result = await db.execute(select(User).where(User.id == user_id).with_for_update().limit(1)) user = result.scalar_one_or_none() if not user: raise ValueError("User not found") return user result = await db.execute(select(User).where(User.id == user_id).with_for_update().limit(1)) user = result.scalar_one_or_none() if not user: raise ValueError("User not found") if biz_key: existing = await _get_existing_credit_record_by_biz_key(db, user_id=user_id, biz_key=biz_key) if existing: return user if float(user.credits or 0) < amount: raise InsufficientCreditsError() user.credits = round(float(user.credits or 0) - amount, 2) meta_kwargs = {} if record_meta: if isinstance(record_meta, CreditRecordMeta): record_meta = await with_user_snapshot(db, record_meta, user_id) meta_kwargs = record_meta.to_record_kwargs() elif isinstance(record_meta, dict): meta_kwargs = {k: v for k, v in record_meta.items() if v is not None} record = CreditRecord( id=generate_id(), user_id=user_id, type="consume", amount=-amount, balance_after=user.credits, description=description, related_id=related_id, biz_key=biz_key, refund_for_biz_key=refund_for_biz_key, **meta_kwargs, ) db.add(record) await db.flush() return user async def add_credits( db: AsyncSession, user_id: str, amount: float, description: str, related_id: str | None = None, *, record_type: str = "recharge", biz_key: str | None = None, refund_for_biz_key: str | None = None, record_meta: CreditRecordMeta | dict | None = None, ) -> User: """增加用户积分,并写入流水。 record_type 默认保持原来的 recharge;生成失败回退时传 refund。 biz_key 不为空时幂等,重复调用不会重复加积分。 """ amount = round(float(amount or 0), 2) result = await db.execute(select(User).where(User.id == user_id).with_for_update().limit(1)) user = result.scalar_one_or_none() if not user: raise ValueError("User not found") if biz_key: existing = await _get_existing_credit_record_by_biz_key(db, user_id=user_id, biz_key=biz_key) if existing: return user if amount <= 0: return user user.credits = round(float(user.credits or 0) + amount, 2) meta_kwargs = {} if record_meta: if isinstance(record_meta, CreditRecordMeta): record_meta = await with_user_snapshot(db, record_meta, user_id) meta_kwargs = record_meta.to_record_kwargs() elif isinstance(record_meta, dict): meta_kwargs = {k: v for k, v in record_meta.items() if v is not None} record = CreditRecord( id=generate_id(), user_id=user_id, type=record_type, amount=amount, balance_after=user.credits, description=description, related_id=related_id, biz_key=biz_key, refund_for_biz_key=refund_for_biz_key, **meta_kwargs, ) db.add(record) await db.flush() return user async def refund_credits( db: AsyncSession, user_id: str, amount: float, description: str, related_id: str | None = None, *, biz_key: str | None = None, refund_for_biz_key: str | None = None, record_meta: CreditRecordMeta | dict | None = None, ) -> User: """生成失败积分回退。""" return await add_credits( db, user_id=user_id, amount=amount, description=description, related_id=related_id, record_type="refund", biz_key=biz_key, refund_for_biz_key=refund_for_biz_key, record_meta=record_meta, ) async def get_records(db: AsyncSession, user_id: str, page: int = 1, page_size: int = 20) -> tuple[list[CreditRecord], int]: count_query = select(func.count(CreditRecord.id)).where(CreditRecord.user_id == user_id) total = (await db.execute(count_query)).scalar() or 0 result = await db.execute( select(CreditRecord) .where(CreditRecord.user_id == user_id) .order_by(CreditRecord.created_at.desc()) .offset((page - 1) * page_size) .limit(page_size) ) return list(result.scalars().all()), total