#!/usr/bin/env python3
"""
baoyu-imagine wrapper for nanobot
简化信息图生成流程
"""

import os
import sys
import subprocess
import tempfile
from pathlib import Path

# 配置
BAOYU_IMAGINE_DIR = "/home/ubuntu/forgithub/baoyu-skills/skills/baoyu-imagine"
DEFAULT_PROVIDER = "seedream"

def get_api_key():
    """获取 API Key"""
    # 优先环境变量
    if "ARK_API_KEY" in os.environ:
        return os.environ["ARK_API_KEY"]
    
    # 其次配置文件
    config_path = Path.home() / ".config" / "baoyu" / "keys.conf"
    if config_path.exists():
        with open(config_path) as f:
            for line in f:
                if line.startswith("ARK_API_KEY="):
                    return line.strip().split("=", 1)[1]
    
    raise ValueError("ARK_API_KEY not found. Please set it in ~/.config/baoyu/keys.conf")

def generate_infographic(system_prompt: str, content_prompt: str, output_path: str):
    """生成信息图"""
    
    api_key = get_api_key()
    
    # 创建临时 prompt 文件
    with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False) as f:
        system_file = f.name
        f.write(system_prompt)
    
    with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False) as f:
        content_file = f.name
        f.write(content_prompt)
    
    try:
        # 调用 baoyu-imagine
        cmd = [
            "npx", "-y", "bun",
            f"{BAOYU_IMAGINE_DIR}/scripts/main.ts",
            "--provider", DEFAULT_PROVIDER,
            "--promptfiles", system_file, content_file,
            "--image", output_path
        ]
        
        env = os.environ.copy()
        env["ARK_API_KEY"] = api_key
        
        result = subprocess.run(
            cmd,
            env=env,
            capture_output=True,
            text=True,
            timeout=300
        )
        
        if result.returncode != 0:
            print(f"Error: {result.stderr}", file=sys.stderr)
            return False
        
        return True
        
    finally:
        # 清理临时文件
        os.unlink(system_file)
        os.unlink(content_file)

if __name__ == "__main__":
    if len(sys.argv) < 4:
        print("Usage: python generate.py <system-prompt.md> <content-prompt.md> <output.png>")
        sys.exit(1)
    
    with open(sys.argv[1]) as f:
        system = f.read()
    
    with open(sys.argv[2]) as f:
        content = f.read()
    
    success = generate_infographic(system, content, sys.argv[3])
    sys.exit(0 if success else 1)