#!/bin/bash
# 图片生成脚本 - 活动发布 skill
# 使用豆包大模型根据活动内容生成提示词，然后生成海报
# 包含：自动添加Logo、保存规范提示词

# 获取脚本所在目录的绝对路径
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SKILL_DIR="$(dirname "$SCRIPT_DIR")"
ASSETS_DIR="$SKILL_DIR/assets"

# 确保 assets 目录存在
mkdir -p "$ASSETS_DIR"

# 加载环境变量
if [ -f "$SCRIPT_DIR/.env" ]; then
    source "$SCRIPT_DIR/.env"
elif [ -f "$SKILL_DIR/.env" ]; then
    source "$SKILL_DIR/.env"
fi

# 脚本路径
PROMPT_SCRIPT="$SCRIPT_DIR/generate_prompt.sh"
LOGO_SCRIPT="$SCRIPT_DIR/add_logo.py"
LOGO_FILE="$ASSETS_DIR/quant-speed.svg"

# 默认值
DOUBAO_API_KEY="${DOUBAO_API_KEY:-db1f8b60-0ffc-473c-98da-40daa3a95df8}"
NANOBANANA_API_KEY="${NANOBANANA_API_KEY:-sk-Cx883TlPrrn0nUvQFmU5wSt99OFv7tiIhg4Re5yHttNtN1Am}"
NANOBANANA_BASE_URL="${NANOBANANA_BASE_URL:-http://zx2.52youxi.cc:3000}"

# 辅助函数：添加Logo
add_logo_to_image() {
    local image_path="$1"
    if [ -f "$LOGO_SCRIPT" ] && [ -f "$LOGO_FILE" ]; then
        echo "   🔄 正在添加Logo..."
        python3 "$LOGO_SCRIPT" "$image_path" "$LOGO_FILE"
    else
        echo "   ⚠️ Logo脚本或文件不存在，跳过添加Logo"
    fi
}

# 辅助函数：保存提示词
save_prompt_info() {
    local image_path="$1"
    local prompt="$2"
    local model="$3"
    local timestamp=$(date +%Y-%m-%dT%H:%M:%S%z)
    
    local txt_file="${image_path}.txt"
    echo "Image: $(basename "$image_path")
Prompt: $prompt
Model: $model
Timestamp: $timestamp" > "$txt_file"
    echo "   💾 提示词已保存: $(basename "$txt_file")"
}

# 使用AI生成提示词
generate_ai_prompts() {
    local event_content="$1"
    
    echo "======================================"
    echo "🤖 步骤1: 使用豆包AI生成海报提示词"
    echo "======================================"
    echo ""
    
    # 调用AI生成提示词
    local result=$(bash "$PROMPT_SCRIPT" ai "$event_content")
    
    # 提取提示词
    local doubao_prompt=$(echo "$result" | grep "DOUBAO_PROMPT=" | sed 's/DOUBAO_PROMPT=//')
    local nano_prompt=$(echo "$result" | grep "NANO_PROMPT=" | sed 's/NANO_PROMPT=//')
    
    if [ -z "$doubao_prompt" ] || [ -z "$nano_prompt" ]; then
        echo "❌ AI提示词生成失败，使用默认提示词"
        doubao_prompt="春节游园会，新春庆典，红灯笼，中国结，舞狮，烟花，喜庆氛围，节日活动"
        nano_prompt="春节，新年，现代简约设计，扁平化风格，渐变色，孟菲斯风格"
    fi
    
    echo "✅ AI提示词生成成功！"
    echo ""
    echo "📝 豆包提示词: ${doubao_prompt:0:80}..."
    echo "📝 创意提示词: ${nano_prompt:0:80}..."
    echo ""
    
    # 返回提示词
    echo "DOUBAO_PROMPT=$doubao_prompt"
    echo "NANO_PROMPT=$nano_prompt"
}

# 生成豆包图片（3张变体）
generate_doubao_images() {
    local base_prompt="$1"
    local event_name="$2"
    
    echo "======================================"
    echo "🎨 步骤2: 生成豆包海报（3张变体）"
    echo "======================================"
    echo ""
    
    # 生成3个不同的提示词变体
    local prompts=$(bash "$PROMPT_SCRIPT" variations "$base_prompt")
    
    local prompt1=$(echo "$prompts" | sed -n '1p')
    local prompt2=$(echo "$prompts" | sed -n '2p')
    local prompt3=$(echo "$prompts" | sed -n '3p')
    
    echo "   提示词变体1: ${prompt1:0:60}..."
    echo "   提示词变体2: ${prompt2:0:60}..."
    echo "   提示词变体3: ${prompt3:0:60}..."
    echo ""
    
    local timestamp=$(date +%Y%m%d_%H%M%S)
    local success_count=0
    
    # 生成3张图片
    for i in 1 2 3; do
        local prompt_var=$(eval echo \$prompt$i)
        local output_file="$ASSETS_DIR/${event_name}_豆包版${i}_${timestamp}.jpg"
        
        echo ">>> 生成第 $i 张..."
        
        response=$(curl -s -X POST "https://ark.cn-beijing.volces.com/api/v3/images/generations" \
            -H "Content-Type: application/json" \
            -H "Authorization: Bearer $DOUBAO_API_KEY" \
            -d "{
                \"model\": \"doubao-seedream-5-0-260128\",
                \"prompt\": \"$prompt_var\",
                \"sequential_image_generation\": \"disabled\",
                \"response_format\": \"url\",
                \"size\": \"2K\",
                \"stream\": false,
                \"watermark\": true
            }")
        
        image_url=$(echo "$response" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d['data'][0]['url'] if d.get('data') else '')" 2>/dev/null)
        
        if [ -n "$image_url" ]; then
            curl -s -o "$output_file" "$image_url"
            echo "   ✅ 第 $i 张已保存: ${output_file}"
            
            # 添加Logo
            add_logo_to_image "$output_file"
            
            # 保存提示词
            save_prompt_info "$output_file" "$prompt_var" "doubao-seedream-5-0-260128"
            
            success_count=$((success_count + 1))
        else
            echo "   ❌ 第 $i 张生成失败"
        fi
    done
    
    echo ""
    echo "✅ 豆包图片生成完成！成功 $success_count/3 张"
}

# 生成 Nanobanana 图片
generate_nanobanana_image() {
    local prompt="$1"
    local event_name="$2"
    
    echo "======================================"
    echo "🎨 步骤3: 生成 Nanobanana 海报（创意版）"
    echo "======================================"
    echo ""
    echo "   提示词: ${prompt:0:60}..."
    echo ""
    
    local timestamp=$(date +%Y%m%d_%H%M%S)
    local output_file="$ASSETS_DIR/${event_name}_创意版_${timestamp}.png"
    
    response=$(curl -s -X POST "${NANOBANANA_BASE_URL}/v1/chat/completions" \
        -H "Content-Type: application/json" \
        -H "Authorization: Bearer $NANOBANANA_API_KEY" \
        -d "{
            \"model\": \"gemini-3-pro-image-preview\",
            \"messages\": [{\"role\": \"user\", \"content\": \"生成一张图片：$prompt\"}],
            \"max_tokens\": 4000
        }")
    
    # 提取base64图片数据
    img_data=$(echo "$response" | python3 -c "
import sys, json, re
import base64
try:
    d = json.load(sys.stdin)
    content = d['choices'][0]['message']['content']
    match = re.search(r'data:image/(\w+);base64,(.+)', content, re.DOTALL)
    if match:
        img_type = match.group(1)
        img_b64 = match.group(2).strip()
        img_bytes = base64.b64decode(img_b64)
        import base64 as b64mod
        print(b64mod.b64encode(img_bytes).decode('utf-8'))
except Exception as e:
    print('')
" 2>/dev/null)
    
    if [ -n "$img_data" ]; then
        echo "$img_data" | base64 -d > "$output_file" 2>/dev/null
        echo "✅ Nanobanana 图片已保存: $output_file"
        
        # 添加Logo
        add_logo_to_image "$output_file"
        
        # 保存提示词
        save_prompt_info "$output_file" "$prompt" "gemini-3-pro-image-preview"
    else
        echo "❌ Nanobanana 图片生成失败"
    fi
}

# 根据活动内容生成所有海报
generate_posters_from_content() {
    local event_content="$1"
    local event_name="$2"
    
    echo "======================================"
    echo "🎉 开始生成活动海报"
    echo "======================================"
    echo ""
    
    # 步骤1: AI生成提示词
    local prompts=$(generate_ai_prompts "$event_content")
    local doubao_prompt=$(echo "$prompts" | grep "DOUBAO_PROMPT=" | sed 's/DOUBAO_PROMPT=//')
    local nano_prompt=$(echo "$prompts" | grep "NANO_PROMPT=" | sed 's/NANO_PROMPT=//')
    
    # 步骤2: 生成豆包图片
    generate_doubao_images "$doubao_prompt" "$event_name"
    echo ""
    
    # 步骤3: 生成Nanobanana图片
    generate_nanobanana_image "$nano_prompt" "$event_name"
    echo ""
    
    echo "======================================"
    echo "✅ 全部海报生成完成！"
    echo "======================================"
    echo ""
    echo "📁 所有图片保存位置: $ASSETS_DIR"
    echo ""
    echo "请选择您喜欢的海报版本："
    echo "  1 - 豆包版1"
    echo "  2 - 豆包版2"
    echo "  3 - 豆包版3"
    echo "  4 - 创意版"
}

# 根据参数调用对应函数
case "$1" in
    ai)
        generate_ai_prompts "$2"
        ;;
    doubao)
        generate_doubao_images "$2" "$3"
        ;;
    nanobanana)
        generate_nanobanana_image "$2" "$3"
        ;;
    all)
        generate_posters_from_content "$2" "$3"
        ;;
    *)
        echo "======================================"
        echo "🖼️  活动海报生成工具"
        echo "======================================"
        echo "用法: $0 all <活动文案> <活动名称>"
        echo ""
        echo "示例: "
        echo '  $0 all "活动名称：春节游园会..." "春节游园会"'
        ;;
esac