#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
测试PosterGenerator组件

该脚本演示如何独立使用PosterGenerator组件创建海报。
"""

import os
import sys
import argparse
import json
import time
import random
from pathlib import Path
from PIL import Image

# 将项目根目录添加到PATH
project_root = str(Path(__file__).parent.parent.absolute())
if project_root not in sys.path:
    sys.path.append(project_root)

from core import poster_gen

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='测试海报生成功能')
    parser.add_argument('--base_dir', '-b', type=str,
                        help='海报素材基础目录路径，包含贴纸/字体/边框等')
    parser.add_argument('--image', '-i', type=str,
                        help='底图路径')
    parser.add_argument('--output_dir', '-o', type=str, default='poster_output',
                        help='输出海报保存目录，默认为"poster_output"')
    parser.add_argument('--config', '-c', type=str,
                        help='海报配置文件路径')
    return parser.parse_args()

def load_config(config_path="poster_gen_config.json"):
    """加载配置文件"""
    if not os.path.exists(config_path):
        config_path = os.path.join(project_root, config_path)
        if not os.path.exists(config_path):
            raise FileNotFoundError(f"配置文件未找到: {config_path}")
    
    with open(config_path, 'r', encoding='utf-8') as f:
        config = json.load(f)
    return config

def create_test_image(output_path, size=(900, 1200)):
    """创建测试底图"""
    # 创建一个渐变背景
    img = Image.new('RGB', size, (0, 0, 0))
    pixels = img.load()
    
    width, height = size
    for x in range(width):
        for y in range(height):
            r = int(255 * x / width)
            g = int(255 * y / height)
            b = int(127 * (x + y) / (width + height))
            pixels[x, y] = (r, g, b)
    
    # 保存图片
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    img.save(output_path)
    print(f"创建测试底图: {output_path}")
    return output_path

def sample_text_data():
    """生成示例文本数据"""
    # 几组示例文本数据，便于测试不同文案效果
    samples = [
        {
            "title": "收藏！泰宁古城",
            "subtitle": "",
            "additional_texts": [
                {"text": "周末游｜必去", "position": "bottom", "size_factor": 0.5},
                {"text": "千年古韵", "position": "bottom", "size_factor": 0.5}
            ]
        },
        {
            "title": "拿捏龙门石窟",
            "subtitle": "",
            "additional_texts": [
                {"text": "亲子研学", "position": "bottom", "size_factor": 0.5},
                {"text": "文化瑰宝", "position": "bottom", "size_factor": 0.5}
            ]
        },
        {
            "title": "厦门鼓浪屿",
            "subtitle": "音乐之岛",
            "additional_texts": [
                {"text": "出片｜必打卡", "position": "bottom", "size_factor": 0.5}
            ]
        }
    ]
    return random.choice(samples)

def main():
    """主函数"""
    args = parse_arguments()
    
    # 确保输出目录存在
    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)
    
    try:
        # 尝试加载系统配置
        try:
            config = load_config()
            print("系统配置加载成功")
        except Exception as e:
            print(f"系统配置加载失败，将使用默认值和命令行参数: {e}")
            config = {}
        
        # 确定海报素材基础目录
        base_dir = args.base_dir
        if not base_dir:
            base_dir = config.get("poster_assets_base_dir")
            if not base_dir:
                print("警告: 未提供海报素材基础目录，将使用默认路径")
                base_dir = "poster_assets"
        
        print(f"使用海报素材基础目录: {base_dir}")
        
        # 初始化PosterGenerator
        print("初始化PosterGenerator...")
        poster_generator = poster_gen.PosterGenerator(base_dir=base_dir)
        
        # 确定底图路径
        image_path = args.image
        if not image_path or not os.path.exists(image_path):
            # 尝试从配置中获取示例图片
            try:
                img_base_dir = config.get("image_base_dir")
                object_name = "测试景点"
                img_subdir = config.get("modify_image_subdir", "modify")
                
                if img_base_dir and os.path.exists(img_base_dir):
                    example_dir = os.path.join(img_base_dir, img_subdir)
                    if os.path.exists(example_dir):
                        for root, dirs, files in os.walk(example_dir):
                            for file in files:
                                if file.lower().endswith(('.jpg', '.jpeg', '.png')):
                                    image_path = os.path.join(root, file)
                                    break
                            if image_path:
                                break
            except Exception:
                pass
                
            # 如果仍未找到图片，创建测试图片
            if not image_path or not os.path.exists(image_path):
                print("未找到有效底图，创建测试图片...")
                image_path = create_test_image(os.path.join(args.output_dir, "test_base_image.jpg"))
        
        print(f"使用底图: {image_path}")
        
        # 加载底图
        try:
            base_image = Image.open(image_path).convert('RGBA')
            print(f"底图加载成功，尺寸: {base_image.size}")
        except Exception as e:
            print(f"底图加载失败，将创建纯色底图: {e}")
            base_image = Image.new('RGBA', (900, 1200), (200, 200, 200, 255))
        
        # 确定文本数据
        text_data = None
        if args.config and os.path.exists(args.config):
            try:
                with open(args.config, 'r', encoding='utf-8') as f:
                    config_data = json.load(f)
                
                # 检查是否是海报配置列表
                if isinstance(config_data, list) and len(config_data) > 0:
                    # 使用第一个配置
                    poster_config = config_data[0]
                    text_data = {
                        "title": poster_config.get("main_title", "默认标题"),
                        "subtitle": "",
                        "additional_texts": []
                    }
                    
                    # 添加文本内容
                    texts = poster_config.get("texts", [])
                    if texts and len(texts) > 0:
                        for i, text in enumerate(texts):
                            text_data["additional_texts"].append({
                                "text": text,
                                "position": "bottom",
                                "size_factor": 0.5
                            })
                
                print(f"从配置文件加载文本数据成功: {text_data}")
            except Exception as e:
                print(f"配置文件解析失败: {e}")
                text_data = None
        
        # 如果未提供配置或解析失败，使用示例文本
        if not text_data:
            text_data = sample_text_data()
            print(f"使用示例文本数据: {text_data}")
        
        # 生成海报
        print("\n开始生成海报...")
        start_time = time.time()
        
        # 设置是否使用文本背景
        os.environ['USE_TEXT_BG'] = 'False'
        
        # 生成海报
        poster_image = poster_generator.create_poster(base_image, text_data)
        
        if poster_image:
            # 保存海报
            output_path = os.path.join(args.output_dir, "generated_poster.jpg")
            poster_image.save(output_path, "JPEG", quality=95)
            
            elapsed = time.time() - start_time
            print(f"海报生成成功，耗时: {elapsed:.2f}秒")
            print(f"海报已保存至: {output_path}")
        else:
            print("海报生成失败")
    
    except Exception as e:
        print(f"程序执行出错: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main() 