ai_v/services/generation_service.py

import json
import threading
import uuid

import requests

from config import Config
from extensions import db, redis_client
from models import GenerationRecord, SystemDict, User
from services.logger import system_logger
from services.task_service import process_image_generation, process_video_generation
from utils import get_api_candidates, get_proxied_url, should_switch_to_backup


def get_model_cost(model_value, is_video=False):
    """获取模型消耗积分。"""
    dict_type = "video_model" if is_video else "ai_model"
    model_dict = SystemDict.query.filter_by(dict_type=dict_type, value=model_value).first()

    if model_dict:
        return model_dict.cost

    # 默认计费兜底
    if is_video:
        return 15 if "pro" in model_value.lower() or "3.1" in model_value else 10
    return 1


def validate_generation_request(user, data):
    """校验生图请求并返回 (api_key, target_api, cost, use_trial, error)。"""
    mode = data.get("mode", "trial")
    is_premium = data.get("is_premium", False)
    input_key = data.get("apiKey")
    model_value = data.get("model")
    is_gemini_flash_image_preview = model_value == "gemini-3.1-flash-image-preview"
    is_gpt_image_edit = model_value == "gpt-image-2"

    if isinstance(is_premium, str):
        is_premium = is_premium.lower() in ("1", "true", "yes", "on")

    # gpt-image-2 统一走 AI_BASE_URL 下的 generations 接口，参考图放在 image 数组中。
    target_api = (
        Config.GPT_IMAGE_GENERATION_API
        if is_gpt_image_edit
        else Config.AI_API
    )
    api_key = None
    use_trial = False

    if mode == "key":
        api_key = input_key or user.api_key
        if not api_key:
            return None, None, 0, False, "请先输入您的 API 密钥"

        # 用户手动输入了新 key 时，同步保存到账号
        if input_key and input_key != user.api_key:
            user.api_key = input_key
            db.session.commit()
    else:
        if user.points <= 0:
            return None, None, 0, False, "可用积分已耗尽，请充值或切换至自定义 Key 模式"

        # 积分模式下，gpt-image-2 跟随 AI_BASE_URL 的内置高级线路。
        if is_gpt_image_edit:
            api_key = Config.PREMIUM_KEY
        elif is_premium:
            api_key = (
                Config.GEMINI_FLASH_IMAGE_PREMIUM_KEY
                if is_gemini_flash_image_preview
                else Config.PREMIUM_KEY
            )
        elif is_gemini_flash_image_preview:
            api_key = Config.PREMIUM_KEY
        else:
            api_key = Config.TRIAL_KEY

        target_api = (
            Config.GPT_IMAGE_GENERATION_API
            if is_gpt_image_edit
            else Config.TRIAL_API
        )
        use_trial = True

    # 计算本次任务消耗
    cost = get_model_cost(model_value, is_video=False)
    if use_trial and is_premium:
        cost *= 2

    if use_trial and user.points < cost:
        return None, None, cost, True, f"可用积分不足，本次需要 {cost} 积分"

    return api_key, target_api, cost, use_trial, None


def deduct_points(user_id, cost):
    """原子扣减积分。"""
    user = db.session.query(User).filter_by(id=user_id).populate_existing().with_for_update().first()
    if user:
        user.points -= cost
        user.has_used_points = True
        db.session.commit()


def refund_points(user_id, cost):
    """原子退回积分。"""
    try:
        user = db.session.query(User).filter_by(id=user_id).populate_existing().with_for_update().first()
        if user:
            user.points += cost
            db.session.commit()
    except Exception:
        db.session.rollback()


def handle_chat_generation_sync(user_id, api_key, model_value, prompt, use_trial, cost):
    """同步处理聊天类模型。"""
    chat_payload = {
        "model": model_value,
        "messages": [{"role": "user", "content": prompt}],
    }
    try:
        resp = None
        last_error = None
        candidates = get_api_candidates(Config.CHAT_API, api_key)

        for idx, candidate in enumerate(candidates):
            headers = {
                "Authorization": f"Bearer {candidate['api_key']}",
                "Content-Type": "application/json",
            }
            try:
                resp = requests.post(
                    get_proxied_url(candidate["url"]),
                    json=chat_payload,
                    headers=headers,
                    timeout=Config.PROXY_TIMEOUT_LONG,
                )
                if resp.status_code == 200:
                    break

                last_error = resp.text
                has_backup = idx < len(candidates) - 1
                if has_backup and should_switch_to_backup(resp.status_code):
                    system_logger.warning(
                        "聊天主线路失败，切换备用线路",
                        user_id=user_id,
                        model=model_value,
                        status_code=resp.status_code,
                    )
                    continue
                break
            except requests.RequestException as e:
                last_error = str(e)
                if idx < len(candidates) - 1:
                    system_logger.warning(
                        "聊天请求异常，切换备用线路",
                        user_id=user_id,
                        model=model_value,
                        error=last_error,
                    )
                    continue
                raise

        if not resp or resp.status_code != 200:
            if use_trial:
                refund_points(user_id, cost)
            return {"error": last_error or "聊天请求失败"}, resp.status_code if resp else 500

        api_result = resp.json()
        content = api_result["choices"][0]["message"]["content"]

        # 非验光单解读的文本生成写入历史
        if prompt != "解读验光单":
            new_record = GenerationRecord(
                user_id=user_id,
                prompt=prompt,
                model=model_value,
                cost=cost,
                image_urls=json.dumps([{"type": "text", "content": content}]),
            )
            db.session.add(new_record)
            db.session.commit()

        return {
            "data": [{"content": content, "type": "text"}],
            "message": "生成成功",
        }, 200
    except Exception as e:
        if use_trial:
            refund_points(user_id, cost)
        return {"error": str(e)}, 500


def start_async_image_task(app, user_id, payload, api_key, target_api, cost, mode, model_value, use_trial=False):
    """启动异步生图任务。"""
    task_id = str(uuid.uuid4())

    log_msg = "用户发起验光单解读" if payload.get("prompt") == "解读验光单" else "用户发起生图任务"
    system_logger.info(log_msg, model=model_value, mode=mode)

    redis_client.setex(
        f"task:{task_id}",
        3600,
        json.dumps({"status": "queued", "message": "任务已提交，等待处理..."}),
    )

    threading.Thread(
        target=process_image_generation,
        args=(app, user_id, task_id, payload, api_key, target_api, cost, use_trial),
    ).start()

    return task_id


def validate_video_request(user, data):
    """校验视频生成请求。"""
    if user.points <= 0:
        return None, 0, "可用积分不足，请先充值"

    model_value = data.get("model", "veo3.1")
    cost = get_model_cost(model_value, is_video=True)

    if user.points < cost:
        return None, cost, f"积分不足，生成该视频需要 {cost} 积分"

    return model_value, cost, None


def start_async_video_task(app, user_id, payload, cost, model_value):
    """启动异步视频任务。"""
    api_key = Config.VIDEO_KEY
    task_id = str(uuid.uuid4())

    system_logger.info("用户发起视频生成任务 (积分模式)", model=model_value, cost=cost)

    redis_client.setex(
        f"task:{task_id}",
        3600,
        json.dumps({"status": "queued", "message": "视频任务已提交，准备开始处理..."}),
    )

    threading.Thread(
        target=process_video_generation,
        args=(app, user_id, task_id, payload, api_key, cost, True),
    ).start()

    return task_id