ai_v/blueprints/api.py

from flask import Blueprint, request, jsonify, session, current_app
from urllib.parse import unquote, urljoin, urlparse
from config import Config
from extensions import db, redis_client, s3_client
from models import User, SavedPrompt
from middlewares.auth import login_required
from services.logger import system_logger
import json

from services.system_service import (
    get_system_config_data,
    get_user_latest_notification,
    mark_notification_as_read,
)
from services.history_service import get_user_history_data
from services.file_service import handle_file_uploads, get_remote_file_stream
from services.generation_service import (
    validate_generation_request,
    deduct_points,
    handle_chat_generation_sync,
    start_async_image_task,
    validate_video_request,
    start_async_video_task,
)
from services.aigc_detector_service import detect_aigc_image

api_bp = Blueprint("api", __name__)


@api_bp.route("/api/task_status/<task_id>")
@login_required
def get_task_status(task_id):
    try:
        data = redis_client.get(f"task:{task_id}")
        if not data:
            return jsonify({"status": "pending"})

        if isinstance(data, bytes):
            data = data.decode("utf-8")

        return jsonify(json.loads(data))
    except Exception as e:
        system_logger.error(f"查询任务状态异常: {str(e)}")
        return jsonify({"status": "error", "message": "状态查询失败"})


@api_bp.route("/api/config")
def get_config():
    try:
        return jsonify(get_system_config_data())
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@api_bp.route("/api/upload", methods=["POST"])
@login_required
def upload():
    try:
        files = request.files.getlist("images")
        img_urls, download_urls = handle_file_uploads(files)
        system_logger.info(f"用户上传文件: {len(files)} 张", user_id=session.get("user_id"))
        return jsonify({"urls": img_urls, "download_urls": download_urls})
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@api_bp.route("/api/generate", methods=["POST"])
@login_required
def generate():
    try:
        user_id = session.get("user_id")
        user = db.session.get(User, user_id)
        data = request.json if request.is_json else request.form

        image_urls = data.get(
            "image",
            data.get("image_urls", data.get("image_data", [])),
        ) or []
        if isinstance(image_urls, str):
            image_urls = [image_urls]
        elif not isinstance(image_urls, (list, tuple)):
            return jsonify({"error": "参考图地址格式无效"}), 400

        normalized_image_urls = []
        for image_url in image_urls:
            if not isinstance(image_url, str) or not image_url.strip():
                return jsonify({"error": "参考图地址无效"}), 400

            image_url = image_url.strip()
            parsed_input = urlparse(image_url)
            if image_url.startswith("/"):
                absolute_url = urljoin(request.url_root, image_url)
            elif parsed_input.scheme in ("http", "https") and parsed_input.netloc:
                absolute_url = image_url
            else:
                return jsonify({"error": "参考图必须使用 HTTP(S) URL，不能使用 Base64"}), 400

            parsed_url = urlparse(absolute_url)
            if parsed_url.scheme not in ("http", "https") or not parsed_url.netloc:
                return jsonify({"error": "参考图必须使用 HTTP(S) URL，不能使用 Base64"}), 400
            normalized_image_urls.append(absolute_url)

        api_key, target_api, cost, use_trial, error = validate_generation_request(user, data)
        if error:
            return jsonify({"error": error}), 400

        if use_trial:
            deduct_points(user_id, cost)

        model_value = data.get("model")
        prompt = (data.get("prompt") or "").strip()
        if not prompt:
            return jsonify({"error": "提示词不能为空"}), 400
        model_lower = model_value.lower()
        is_chat_model = ("gemini" in model_lower or "gpt" in model_lower) and "image" not in model_lower

        if is_chat_model:
            result, status_code = handle_chat_generation_sync(
                user_id, api_key, model_value, prompt, use_trial, cost
            )
            return jsonify(result), status_code

        if model_value == "gpt-image-2":
            payload = {
                "prompt": prompt,
                "model": model_value,
                "response_format": "url",
                "size": data.get("size") or "2048x2048",
                "quality": "high",
            }
            if normalized_image_urls:
                payload["image"] = normalized_image_urls
        else:
            payload = {
                "prompt": prompt,
                "model": model_value,
                "response_format": "url",
                "aspect_ratio": data.get("ratio"),
            }
            if normalized_image_urls:
                payload["image"] = normalized_image_urls

            if model_value in ("nano-banana-2", "gemini-3.1-flash-image-preview") and data.get("size"):
                payload["image_size"] = data.get("size")

        app = current_app._get_current_object()
        task_id = start_async_image_task(
            app,
            user_id,
            payload,
            api_key,
            target_api,
            cost,
            data.get("mode"),
            model_value,
            use_trial,
        )

        return jsonify({
            "task_id": task_id,
            "message": "已开启异步生成任务",
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@api_bp.route("/api/aigc-detect", methods=["POST"])
@login_required
def aigc_detect():
    try:
        user_id = session.get("user_id")
        data = request.json if request.is_json else request.form
        result, status_code = detect_aigc_image(
            user_id,
            (data.get("image_url") or data.get("url") or "").strip(),
            request.url_root,
        )
        return jsonify(result), status_code
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@api_bp.route("/api/video/generate", methods=["POST"])
@login_required
def video_generate():
    try:
        user_id = session.get("user_id")
        user = db.session.get(User, user_id)
        data = request.json

        model_value, cost, error = validate_video_request(user, data)
        if error:
            return jsonify({"error": error}), 400

        deduct_points(user_id, cost)

        payload = {
            "model": model_value,
            "prompt": data.get("prompt"),
            "enhance_prompt": data.get("enhance_prompt", False),
            "images": data.get("images", []),
            "aspect_ratio": data.get("aspect_ratio", "9:16"),
        }

        if payload.get("images"):
            payload["images"] = [img.split(",", 1)[1] if "," in img else img for img in payload["images"]]

        app = current_app._get_current_object()
        task_id = start_async_video_task(app, user_id, payload, cost, model_value)

        return jsonify({
            "task_id": task_id,
            "message": "视频生成任务已提交，系统正在处理中",
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@api_bp.route("/api/notifications/latest", methods=["GET"])
@login_required
def get_latest_notification():
    try:
        user_id = session.get("user_id")
        data = get_user_latest_notification(user_id)
        return jsonify(data)
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@api_bp.route("/api/notifications/read", methods=["POST"])
@login_required
def mark_notif_read():
    try:
        user_id = session.get("user_id")
        data = request.json
        notif_id = data.get("id")
        if not notif_id:
            return jsonify({"error": "缺少通知 ID"}), 400

        mark_notification_as_read(user_id, notif_id)
        return jsonify({"status": "ok"})
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@api_bp.route("/api/history", methods=["GET"])
@login_required
def get_history():
    try:
        user_id = session.get("user_id")
        page = request.args.get("page", 1, type=int)
        per_page = request.args.get("per_page", 10, type=int)
        filter_type = request.args.get("filter_type", "all")
        data = get_user_history_data(user_id, page, per_page, filter_type)
        return jsonify(data)
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@api_bp.route("/api/stats/points", methods=["GET"])
@login_required
def get_user_point_stats():
    from services.stats_service import get_point_stats

    user_id = session.get("user_id")
    days = request.args.get("days", 7, type=int)
    return jsonify(get_point_stats(user_id, days))


@api_bp.route("/api/stats/details", methods=["GET"])
@login_required
def get_user_point_details():
    from services.stats_service import get_point_details

    user_id = session.get("user_id")
    page = request.args.get("page", 1, type=int)
    return jsonify(get_point_details(user_id, page))


@api_bp.route("/api/download_proxy", methods=["GET"])
@login_required
def download_proxy():
    import time

    url = request.args.get("url")
    default_name = f"video-{int(time.time())}.mp4"
    filename = request.args.get("filename") or default_name

    if not url:
        return jsonify({"error": "缺少 URL 参数"}), 400

    try:
        parsed_url = urlparse(url)
        public_prefix = Config.MINIO["public_url"].rstrip("/") + "/"
        is_same_origin_file = (
            parsed_url.path.startswith(public_prefix)
            and (not parsed_url.scheme or parsed_url.netloc == request.host)
        )
        if is_same_origin_file:
            object_key = unquote(parsed_url.path[len(public_prefix):])
            file_obj = s3_client.get_object(Bucket=Config.MINIO["bucket"], Key=object_key)
            headers = {
                "Content-Type": file_obj.get("ContentType") or "application/octet-stream",
                "Content-Disposition": f'attachment; filename="{filename}"',
                "Cache-Control": "private, max-age=0",
            }

            def generate_minio_stream():
                try:
                    for chunk in file_obj["Body"].iter_chunks(chunk_size=4096):
                        yield chunk
                finally:
                    file_obj["Body"].close()

            return current_app.response_class(generate_minio_stream(), headers=headers)

        req, headers = get_remote_file_stream(url)
        headers["Content-Disposition"] = f'attachment; filename="{filename}"'

        def generate_stream():
            try:
                for chunk in req.iter_content(chunk_size=4096):
                    yield chunk
            finally:
                req.close()

        return current_app.response_class(generate_stream(), headers=headers)
    except Exception as e:
        system_logger.error(f"代理下载失败: {str(e)}")
        return jsonify({"error": "下载失败"}), 500


@api_bp.route("/api/prompts", methods=["GET"])
@login_required
def list_prompts():
    try:
        user_id = session.get("user_id")
        prompts = SavedPrompt.query.filter_by(user_id=user_id).order_by(SavedPrompt.created_at.desc()).all()
        return jsonify([
            {
                "id": p.id,
                "label": p.title,
                "value": p.prompt,
            }
            for p in prompts
        ])
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@api_bp.route("/api/prompts", methods=["POST"])
@login_required
def save_prompt():
    try:
        user_id = session.get("user_id")
        data = request.json
        title = data.get("title")
        prompt_text = data.get("prompt")

        if not title or not prompt_text:
            return jsonify({"error": "标题和内容不能为空"}), 400

        count = SavedPrompt.query.filter_by(user_id=user_id).count()
        if count >= 50:
            return jsonify({"error": "收藏数量已达上限 (50)"}), 400

        prompt = SavedPrompt(user_id=user_id, title=title, prompt=prompt_text)
        db.session.add(prompt)
        db.session.commit()

        return jsonify({"message": "保存成功", "id": prompt.id})
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@api_bp.route("/api/prompts/<int:id>", methods=["DELETE"])
@login_required
def delete_prompt(id):
    try:
        user_id = session.get("user_id")
        prompt = SavedPrompt.query.filter_by(id=id, user_id=user_id).first()
        if prompt:
            db.session.delete(prompt)
            db.session.commit()
        return jsonify({"message": "删除成功"})
    except Exception as e:
        return jsonify({"error": str(e)}), 500