tw2/qingjian_cs/app/auto_draw.py

from __future__ import annotations

import asyncio
import os
import sys
import tempfile
import uuid
from pathlib import Path
from typing import Any

import requests
from dotenv import load_dotenv

from .config import AUTO_DRAW_ENDPOINT, AUTO_DRAW_TIMEOUT_SECONDS


def _add_legacy_tw_path() -> None:
    root = os.getenv("LEGACY_TW_ROOT", r"D:\main\sandbox\tw_terminator").strip()
    if not root:
        return
    p = Path(root)
    # 先加载 legacy 项目的 .env，确保 service_tuhui_upload 在 import 时拿到正确账号配置
    legacy_env = p / ".env"
    if legacy_env.exists():
        load_dotenv(legacy_env, override=True)
    if p.exists() and str(p) not in sys.path:
        sys.path.insert(0, str(p))


async def _draw_via_legacy_tw(
    image_url: str,
    customer_id: str,
    requirement: str = "",
) -> dict[str, Any]:
    # 优先使用当前项目中拷贝过来的 service_gemini
    from services.service_gemini import GeminiExtractV2Service  # type: ignore

    # 上传模块暂时仍走 legacy（你后续可替换为新项目本地上传实现）
    _add_legacy_tw_path()
    from services.service_tuhui_upload import upload_to_tuhui  # type: ignore

    prompt = requirement.strip() or "按原图做高清修复，保留主体细节，输出清晰可用版本"

    # 1) 下载原图到本地临时文件
    input_path = os.path.join(tempfile.gettempdir(), f"qjcs_in_{uuid.uuid4().hex}.jpg")
    output_path = os.path.join(tempfile.gettempdir(), f"qjcs_out_{uuid.uuid4().hex}.jpg")
    headers = {
        "User-Agent": (
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
            "AppleWebKit/537.36 (KHTML, like Gecko) "
            "Chrome/122.0.0.0 Safari/537.36"
        ),
        "Referer": "https://www.taobao.com/",
        "Accept": "image/avif,image/webp,image/apng,image/*,*/*;q=0.8",
    }
    resp = requests.get(image_url, headers=headers, timeout=AUTO_DRAW_TIMEOUT_SECONDS)
    if resp.status_code != 200:
        return {"ok": False, "error": f"download_http_{resp.status_code}"}
    with open(input_path, "wb") as f:
        f.write(resp.content)

    # 2) 直调你原来的 service_gemini 作图 API
    service = GeminiExtractV2Service()
    ok_extract, msg_extract, _ = await service.extract_pattern(
        input_path=input_path,
        output_path=output_path,
        custom_prompt=prompt,
        aspect_ratio="1:1",
    )
    if not ok_extract:
        return {"ok": False, "error": f"extract_failed:{msg_extract}"}
    if not os.path.exists(output_path):
        return {"ok": False, "error": "extract_no_output_file"}

    # 3) 上传图绘，返回可外发 URL
    ok, link, _ = await upload_to_tuhui(
        output_path,
        title=f"客户{customer_id[-4:]}-预览图" if customer_id else "预览图",
        description="AI自动作图预览",
        price=1,
    )
    if not ok:
        return {"ok": False, "error": str(link)}
    return {"ok": True, "url": str(link)}


def _draw_via_http_endpoint(image_url: str, customer_id: str, requirement: str = "") -> dict[str, Any]:
    if not AUTO_DRAW_ENDPOINT:
        return {"ok": False, "error": "AUTO_DRAW_ENDPOINT not configured"}
    payload = {
        "image_url": image_url,
        "customer_id": customer_id,
        "requirement": requirement,
    }
    resp = requests.post(AUTO_DRAW_ENDPOINT, json=payload, timeout=AUTO_DRAW_TIMEOUT_SECONDS)
    if resp.status_code != 200:
        return {"ok": False, "error": f"http_{resp.status_code}:{resp.text[:200]}"}
    data = resp.json() if resp.text else {}
    url = str(data.get("url", "") or data.get("preview_url", "") or "")
    if not url:
        return {"ok": False, "error": "missing_preview_url"}
    return {"ok": True, "url": url}


async def auto_draw_preview(
    image_url: str,
    customer_id: str,
    requirement: str = "",
) -> dict[str, Any]:
    """
    统一自动作图入口：
    1) 优先走 tw_terminator 的 service_gemini 直调链路
    2) 失败时回退 AUTO_DRAW_ENDPOINT
    """
    try:
        return await _draw_via_legacy_tw(image_url=image_url, customer_id=customer_id, requirement=requirement)
    except Exception as e:
        legacy_error = str(e)

    try:
        data = await asyncio.to_thread(
            _draw_via_http_endpoint,
            image_url,
            customer_id,
            requirement,
        )
        if data.get("ok"):
            return data
        return {"ok": False, "error": f"legacy:{legacy_error}; endpoint:{data.get('error','unknown')}"}
    except Exception as e:
        return {"ok": False, "error": f"legacy:{legacy_error}; endpoint:{e}"}