#!/usr/bin/env python3
"""
AI 配图生成器 - 使用 sensenova API 生成文章配图
支持 OpenAI 兼容格式的图片生成 API
"""

import os, sys, json, datetime, logging, requests, base64
from pathlib import Path

PROJECT_ROOT = Path('/root/openclaw-workspace/projects/yu-zhi-ran')
sys.path.insert(0, str(PROJECT_ROOT))
sys.path.insert(0, str(PROJECT_ROOT / 'platform' / 'backend'))

from dotenv import load_dotenv
load_dotenv(str(PROJECT_ROOT / 'platform' / 'backend' / '.env'))

LOGS_DIR = PROJECT_ROOT / "automation" / "logs"
TODAY = datetime.datetime.now().strftime("%Y-%m-%d")

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(LOGS_DIR / f"ai_image_{TODAY}.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

SENSENOVA_API_KEY = os.getenv("SENSENOVA_API_KEY", "sk-2Bbcf8pSTSl1x2BV5fKtDsUIGdfjKX7M")
SENSENOVA_BASE_URL = os.getenv("SENSENOVA_BASE_URL", "https://token.sensenova.cn/v1")
SENSENOVA_MODEL = os.getenv("SENSENOVA_IMAGE_MODEL", "sensenova-u1-fast")

# Platform image size mapping
PLATFORM_SIZES = {
    "zhihu": "1760x2368",
    "wechat": "1760x2368",
    "xiaohongshu": "1664x2496",
}

def generate_image(prompt: str, platform: str = "zhihu", size: str = None) -> str:
    if not size:
        size = PLATFORM_SIZES.get(platform, "2048x2048")
    headers = {
        "Authorization": f"Bearer {SENSENOVA_API_KEY}",
        "Content-Type": "application/json",
    }
    payload = {
        "model": SENSENOVA_MODEL,
        "prompt": prompt,
        "n": 1,
        "size": size,
        "response_format": "b64_json",
    }
    try:
        resp = requests.post(
            f"{SENSENOVA_BASE_URL}/images/generations",
            headers=headers,
            json=payload,
            timeout=120,
        )
        if resp.status_code == 200:
            data = resp.json()
            if "data" in data and len(data["data"]) > 0:
                item = data["data"][0]
                if "b64_json" in item:
                    return item["b64_json"]
                if "url" in item:
                    img_resp = requests.get(item["url"], timeout=60)
                    if img_resp.status_code == 200:
                        return base64.b64encode(img_resp.content).decode("ascii")
            logger.error(f"API 返回格式异常: {json.dumps(data, ensure_ascii=False)[:500]}")
        else:
            logger.error(f"API 请求失败: {resp.status_code} {resp.text[:500]}")
    except requests.Timeout:
        logger.error("API 请求超时")
    except Exception as e:
        logger.error(f"生成图片失败: {e}")
    return None

def generate_article_images(topic_id: str, title: str, field: str = "", platforms: list = None) -> dict:
    if platforms is None:
        platforms = ["zhihu", "wechat", "xiaohongshu"]
    output_dir = PROJECT_ROOT / "automation" / "images" / "generated" / TODAY
    output_dir.mkdir(parents=True, exist_ok=True)
    results = {}
    for platform in platforms:
        platform_results = {}
        for img_type, prompt_template in [
            ("cover", f"为文章《{title}》生成一张高质量的封面配图，领域：{field}。风格：简约现代，色彩柔和专业，适合{platform}平台。不要文字。"),
            ("illustration", f"为文章《{title}》（领域：{field}）生成一张内容插图，表现核心概念。风格：清新自然，适合{platform}平台。"),
        ]:
            img_b64 = generate_image(prompt_template, platform)
            if img_b64:
                filename = f"{img_type}_{platform}_{topic_id}.png"
                filepath = output_dir / filename
                with open(filepath, "wb") as f:
                    f.write(base64.b64decode(img_b64))
                platform_results[img_type] = str(filepath)
                logger.info(f"[{platform}] {img_type} 生成成功: {filename}")
            else:
                logger.warning(f"[{platform}] {img_type} 生成失败，跳过")
        results[platform] = platform_results
    return results

def save_images_to_db(topic_id: str, platform: str, images: dict):
    try:
        from app.database import SessionLocal
        from app.models import Article
        db = SessionLocal()
        article_id = f"{platform}_{topic_id}"
        article = db.query(Article).filter(Article.id == article_id).first()
        if article:
            existing = article.images or {}
            existing.update(images)
            article.images = existing
            db.commit()
            logger.info(f"[{platform}] 图片路径已保存到 DB")
        db.close()
    except Exception as e:
        logger.error(f"保存图片路径到 DB 失败: {e}")

def main():
    import argparse
    parser = argparse.ArgumentParser(description="AI 配图生成器")
    parser.add_argument("--topic-id", required=True, help="选题ID")
    parser.add_argument("--title", help="文章标题")
    parser.add_argument("--field", default="", help="领域")
    args = parser.parse_args()
    from db_helper import get_topic_by_id
    topic = get_topic_by_id(args.topic_id)
    if not topic:
        print(json.dumps({"ok": False, "error": f"Topic {args.topic_id} not found"}))
        sys.exit(1)
    title = args.title or topic.get("title", "")
    field = args.field or topic.get("field", "")
    logger.info(f"为选题 {args.topic_id} 《{title}》生成配图...")
    results = generate_article_images(args.topic_id, title, field)
    for platform, imgs in results.items():
        if imgs:
            save_images_to_db(args.topic_id, platform, imgs)
    output = {"ok": True, "topic_id": args.topic_id, "images": results}
    print(json.dumps(output, ensure_ascii=False))

if __name__ == "__main__":
    main()