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

import os
import sys
import logging
import argparse
import json
import random
from pathlib import Path
from datetime import datetime
from PIL import Image

# 添加项目根目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from utils.output_handler import FileSystemOutputHandler
from utils.poster_notes_creator import select_additional_images

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def create_sample_metadata(image_dir, num_images=2):
    """创建示例海报元数据
    
    Args:
        image_dir: 图像目录
        num_images: 使用的图像数量
        
    Returns:
        dict: 海报元数据
    """
    # 获取目录中的图像
    image_extensions = ('.jpg', '.jpeg', '.png', '.bmp')
    available_images = [
        f for f in os.listdir(image_dir)
        if os.path.isfile(os.path.join(image_dir, f)) and
        f.lower().endswith(image_extensions)
    ]
    
    if len(available_images) < num_images:
        logger.warning(f"可用图像数量({len(available_images)})少于请求数量({num_images})，将使用所有可用图像")
        selected_images = available_images
    else:
        selected_images = random.sample(available_images, num_images)
    
    # 创建元数据
    metadata = {
        "title": "示例海报",
        "description": "这是一个用于测试额外配图功能的示例海报",
        "collage_images": selected_images,
        "generation_time": datetime.now().isoformat(),
        "style": "modern"
    }
    
    return metadata

def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='测试额外配图功能')
    parser.add_argument('--image_dir', type=str, required=True, 
                        help='包含图像的目录路径')
    parser.add_argument('--output_dir', type=str, default='./output',
                       help='输出目录路径')
    parser.add_argument('--num_images', type=int, default=3,
                       help='要选择的额外配图数量')
    
    args = parser.parse_args()
    
    # 验证图像目录
    if not os.path.exists(args.image_dir) or not os.path.isdir(args.image_dir):
        logger.error(f"图像目录不存在: {args.image_dir}")
        return 1
    
    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)
    
    # 初始化输出处理器
    output_handler = FileSystemOutputHandler(args.output_dir)
    
    # 创建测试运行ID和主题/变体索引
    run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
    topic_index = 1
    variant_index = 1
    
    # 创建一个模拟的海报元数据文件
    metadata = create_sample_metadata(args.image_dir, 2)
    metadata_path = os.path.join(args.output_dir, f"{run_id}_{topic_index}_{variant_index}_metadata.json")
    
    # 保存元数据到文件
    with open(metadata_path, 'w', encoding='utf-8') as f:
        json.dump(metadata, f, ensure_ascii=False, indent=2)
    logger.info(f"模拟海报元数据已保存: {metadata_path}")
    
    # 打印将要使用的"已用"图像
    logger.info(f"模拟海报使用的图像: {', '.join(metadata['collage_images'])}")
    
    # 调用额外配图选择函数
    logger.info(f"开始选择额外配图...")
    image_paths = select_additional_images(
        run_id=run_id,
        topic_index=topic_index,
        variant_index=variant_index,
        poster_metadata_path=metadata_path,
        source_image_dir=args.image_dir,
        num_additional_images=args.num_images,
        output_handler=output_handler
    )
    
    # 打印结果
    if image_paths:
        logger.info(f"已选择并保存 {len(image_paths)} 张额外配图:")
        for i, path in enumerate(image_paths):
            logger.info(f"  {i+1}. {path}")
        
        logger.info(f"额外配图已保存到: {os.path.join(args.output_dir, run_id, f'{topic_index}_{variant_index}', 'image')}")
        logger.info(f"测试成功完成!")
    else:
        logger.error("未能选择任何额外配图.")
        return 1
    
    return 0
    
if __name__ == "__main__":
    sys.exit(main())