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

import os
import sys
import json
import logging
import argparse
from datetime import datetime

# 添加项目根目录到路径
sys.path.append(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 main():
    parser = argparse.ArgumentParser(description='测试额外配图选择功能')
    parser.add_argument('--run_id', type=str, required=True, help='指定运行ID')
    parser.add_argument('--topic_index', type=int, required=True, help='指定主题索引')
    parser.add_argument('--variant_index', type=int, required=True, help='指定变体索引')
    parser.add_argument('--source_dir', type=str, required=True, help='源图像目录路径')
    parser.add_argument('--output_dir', type=str, default='./result', help='输出目录路径')
    parser.add_argument('--num_images', type=int, default=3, help='要选择的配图数量')
    args = parser.parse_args()

    # 检查源目录是否存在
    if not os.path.exists(args.source_dir) or not os.path.isdir(args.source_dir):
        logger.error(f"源图像目录不存在: {args.source_dir}")
        return 1

    # 检查输出目录结构
    run_dir = os.path.join(args.output_dir, args.run_id)
    variant_dir = os.path.join(run_dir, f"{args.topic_index}_{args.variant_index}")
    poster_dir = os.path.join(variant_dir, "poster")
    
    if not os.path.exists(poster_dir):
        logger.error(f"海报目录不存在: {poster_dir}")
        return 1
    
    # 查找海报元数据文件
    metadata_files = [f for f in os.listdir(poster_dir) 
                     if f.endswith("_metadata.json") and os.path.isfile(os.path.join(poster_dir, f))]
    
    if not metadata_files:
        logger.error(f"未找到海报元数据文件")
        return 1
    
    metadata_path = os.path.join(poster_dir, metadata_files[0])
    logger.info(f"使用元数据文件: {metadata_path}")
    
    # 初始化输出处理器
    output_handler = FileSystemOutputHandler(args.output_dir)
    
    # 调用额外配图选择函数
    logger.info(f"开始为运行ID {args.run_id} 主题 {args.topic_index} 变体 {args.variant_index} 选择额外配图...")
    image_paths = select_additional_images(
        run_id=args.run_id,
        topic_index=args.topic_index,
        variant_index=args.variant_index,
        poster_metadata_path=metadata_path,
        source_image_dir=args.source_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, args.run_id, f'{args.topic_index}_{args.variant_index}', 'image')}")
        logger.info("测试成功完成!")
    else:
        logger.error("未能选择任何额外配图")
        return 1
    
    return 0

if __name__ == "__main__":
    sys.exit(main())