how/data/run_batch.py

"""
批量测试脚本：处理作者历史帖子目录下的所有帖子（并发版本，带历史帖子）

功能：
1. 加载最新的20个帖子（按照publish_timestamp从新到旧排序）
2. 为每个帖子加载历史帖子（比当前帖子早的最近15篇）
3. 并发处理所有帖子
"""

import json
import sys
import os
import argparse
from pathlib import Path
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading

# 添加项目根目录到路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

# 手动加载.env文件
def load_env_file(env_path):
    """手动加载.env文件"""
    if not env_path.exists():
        return False

    with open(env_path, 'r') as f:
        for line in f:
            line = line.strip()
            # 跳过注释和空行
            if not line or line.startswith('#'):
                continue
            # 解析KEY=VALUE
            if '=' in line:
                key, value = line.split('=', 1)
                os.environ[key.strip()] = value.strip()

    return True

env_path = project_root / ".env"
if load_env_file(env_path):
    print(f"✅ 已加载环境变量从: {env_path}")
    # 验证API密钥
    api_key = os.environ.get("GEMINI_API_KEY", "")
    if api_key:
        print(f"   GEMINI_API_KEY: {api_key[:10]}...")
else:
    print(f"⚠️  未找到.env文件: {env_path}")

from src.workflows.what_deconstruction_workflow import WhatDeconstructionWorkflow
from src.utils.logger import get_logger

logger = get_logger(__name__)

# 线程安全的输出锁
print_lock = threading.Lock()

# 线程安全的计数器
class ThreadSafeCounter:
    def __init__(self):
        self._lock = threading.Lock()
        self._success = 0
        self._fail = 0

    def increment_success(self):
        with self._lock:
            self._success += 1

    def increment_fail(self):
        with self._lock:
            self._fail += 1

    @property
    def success(self):
        with self._lock:
            return self._success

    @property
    def fail(self):
        with self._lock:
            return self._fail

def safe_print(*args, **kwargs):
    """线程安全的打印函数"""
    with print_lock:
        print(*args, **kwargs)


def load_historical_posts(history_dir, target_timestamp=None, target_post_id=None, max_count=15):
    """
    加载历史帖子（根据publish_timestamp从新到旧排序）

    选择比目标帖子早发布，并且是最近发布的帖子，排除目标帖子本身

    Args:
        history_dir: 历史帖子目录
        target_timestamp: 目标帖子的发布时间戳（可选）
        target_post_id: 目标帖子的ID（用于过滤重复，可选）
        max_count: 最多加载的帖子数量

    Returns:
        list: 历史帖子列表（从新到旧排序）
    """
    history_path = Path(history_dir)

    if not history_path.exists():
        safe_print(f"⚠️  历史帖子目录不存在: {history_path}")
        return []

    # 获取所有JSON文件
    json_files = list(history_path.glob("*.json"))

    if not json_files:
        safe_print(f"⚠️  未找到历史帖子文件")
        return []

    # 读取所有帖子并提取publish_timestamp
    posts_with_timestamp = []
    for file_path in json_files:
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                post_data = json.load(f)

            # 获取发布时间戳，如果不存在则使用0
            timestamp = post_data.get("publish_timestamp", 0)
            post_id = post_data.get("channel_content_id", "")

            posts_with_timestamp.append({
                "file_path": file_path,
                "post_data": post_data,
                "timestamp": timestamp,
                "post_id": post_id
            })
        except Exception as e:
            safe_print(f"   ⚠️  读取文件失败 {file_path.name}: {e}")
            continue

    if not posts_with_timestamp:
        safe_print(f"⚠️  没有成功读取到任何帖子")
        return []

    # 过滤掉目标帖子本身
    if target_post_id is not None:
        posts_with_timestamp = [
            post for post in posts_with_timestamp
            if post["post_id"] != target_post_id
        ]

    # 如果提供了目标时间戳，只保留比目标帖子早的帖子
    if target_timestamp is not None:
        posts_with_timestamp = [
            post for post in posts_with_timestamp
            if post["timestamp"] < target_timestamp
        ]

        if not posts_with_timestamp:
            return []

    # 按照publish_timestamp排序（从新到旧）
    posts_with_timestamp.sort(key=lambda x: x["timestamp"], reverse=True)

    # 选择最近的N篇（从新到旧）
    selected_posts = posts_with_timestamp[:max_count] if len(posts_with_timestamp) > max_count else posts_with_timestamp

    historical_posts = []
    for post_info in selected_posts:
        post_data = post_info["post_data"]

        # 转换为需要的格式
        historical_post = {
            "text": {
                "title": post_data.get("title", ""),
                "body": post_data.get("body_text", ""),
                "hashtags": ""
            },
            "images": post_data.get("images", [])
        }

        historical_posts.append(historical_post)

    return historical_posts


def load_post_files(directory, max_count=20):
    """
    加载作者历史帖子目录下的所有JSON文件，按照publish_timestamp从新到旧排序，取最新的max_count个

    Args:
        directory: 帖子目录
        max_count: 最多加载的帖子数量（默认20）

    Returns:
        list: 帖子文件路径列表（按时间从新到旧排序）
    """
    post_dir = Path(directory)

    if not post_dir.exists():
        raise FileNotFoundError(f"目录不存在: {post_dir}")

    # 获取所有JSON文件
    json_files = list(post_dir.glob("*.json"))

    if not json_files:
        raise FileNotFoundError(f"目录中没有找到JSON文件: {post_dir}")

    # 读取所有帖子并提取publish_timestamp
    posts_with_timestamp = []
    for file_path in json_files:
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                post_data = json.load(f)

            # 获取发布时间戳，如果不存在则使用0
            timestamp = post_data.get("publish_timestamp", 0)

            posts_with_timestamp.append({
                "file_path": file_path,
                "timestamp": timestamp,
                "post_data": post_data
            })
        except Exception as e:
            print(f"⚠️  读取文件失败 {file_path.name}: {e}")
            continue

    if not posts_with_timestamp:
        raise FileNotFoundError(f"没有成功读取到任何帖子")

    # 按照publish_timestamp排序（从新到旧）
    posts_with_timestamp.sort(key=lambda x: x["timestamp"], reverse=True)

    # 取最新的max_count个
    selected_posts = posts_with_timestamp[:max_count]

    print(f"📊 按时间排序并选择最新 {len(selected_posts)} 个帖子:")
    for idx, post_info in enumerate(selected_posts, 1):
        post_data = post_info["post_data"]
        publish_time = post_data.get("publish_time", "未知时间")
        title = post_data.get("title", "无标题")
        print(f"   {idx}. {post_info['file_path'].name}")
        print(f"      标题: {title}")
        print(f"      发布时间: {publish_time}")

    return [post_info["file_path"] for post_info in selected_posts]


def convert_to_workflow_input(raw_data, historical_posts=None):
    """
    将原始数据转换为工作流输入格式

    Args:
        raw_data: 原始帖子数据
        historical_posts: 历史帖子列表（可选）
    """
    # 转换为工作流需要的格式
    images = raw_data.get("images", [])

    input_data = {
        "multimedia_content": {
            "images": images,
            "video": raw_data.get("video", {}),
            "text": {
                "title": raw_data.get("title", ""),
                "body": raw_data.get("body_text", ""),
                "hashtags": ""
            }
        },
        "comments": raw_data.get("comments", []),  # 包含评论数据
        "creator_info": {
            "nickname": raw_data.get("channel_account_name", ""),
            "account_id": raw_data.get("channel_account_id", "")
        }
    }

    # 如果有历史帖子，添加到输入数据中
    if historical_posts:
        input_data["historical_posts"] = historical_posts

    return input_data


def process_single_post(post_file, posts_dir, output_dir, counter, total_count, current_index):
    """处理单个帖子文件（线程安全版本，带历史帖子）"""
    post_name = post_file.stem  # 获取文件名（不含扩展名）
    thread_id = threading.current_thread().name

    safe_print(f"\n{'='*80}")
    safe_print(f"[线程:{thread_id}] 处理帖子: {post_name}")
    safe_print(f"进度: [{current_index}/{total_count}]")
    safe_print(f"{'='*80}")

    # 1. 加载帖子数据
    safe_print(f"\n[线程:{thread_id}][1] 加载帖子数据...")
    try:
        with open(post_file, "r", encoding="utf-8") as f:
            raw_data = json.load(f)

        target_timestamp = raw_data.get('publish_timestamp')
        target_post_id = raw_data.get('channel_content_id')

        safe_print(f"✅ [{thread_id}] 成功加载帖子数据")
        safe_print(f"   - 标题: {raw_data.get('title', 'N/A')}")
        safe_print(f"   - 帖子ID: {target_post_id}")
        safe_print(f"   - 发布时间: {raw_data.get('publish_time', '未知时间')}")
        safe_print(f"   - 图片数: {len(raw_data.get('images', []))}")
        safe_print(f"   - 点赞数: {raw_data.get('like_count', 0)}")
        safe_print(f"   - 评论数: {raw_data.get('comment_count', 0)}")
    except Exception as e:
        safe_print(f"❌ [{thread_id}] 加载帖子数据失败: {e}")
        counter.increment_fail()
        return False

    # 2. 加载历史帖子
    safe_print(f"\n[线程:{thread_id}][2] 加载历史帖子...")
    historical_posts = load_historical_posts(
        posts_dir,
        target_timestamp=target_timestamp,
        target_post_id=target_post_id,
        max_count=15
    )

    if historical_posts:
        safe_print(f"✅ [{thread_id}] 成功加载 {len(historical_posts)} 篇历史帖子")
    else:
        safe_print(f"⚠️  [{thread_id}] 未加载到历史帖子，将使用常规分析模式")

    # 3. 转换数据格式
    safe_print(f"\n[线程:{thread_id}][3] 转换数据格式...")
    try:
        input_data = convert_to_workflow_input(raw_data, historical_posts)
        safe_print(f"✅ [{thread_id}] 数据格式转换成功")
        safe_print(f"   - 历史帖子数: {len(input_data.get('historical_posts', []))}")
    except Exception as e:
        safe_print(f"❌ [{thread_id}] 数据格式转换失败: {e}")
        counter.increment_fail()
        return False

    # 4. 初始化工作流（每个线程创建独立实例，确保线程安全）
    safe_print(f"\n[线程:{thread_id}][4] 初始化工作流...")
    try:
        workflow = WhatDeconstructionWorkflow(
            model_provider="google_genai",
            max_depth=10
        )
        safe_print(f"✅ [{thread_id}] 工作流初始化成功")
    except Exception as e:
        safe_print(f"❌ [{thread_id}] 工作流初始化失败: {e}")
        counter.increment_fail()
        return False

    # 5. 执行工作流
    safe_print(f"\n[线程:{thread_id}][5] 执行工作流...")
    safe_print(f"   注意：这可能需要几分钟时间...")
    try:
        result = workflow.invoke(input_data)
        safe_print(f"✅ [{thread_id}] 工作流执行成功")
    except Exception as e:
        safe_print(f"❌ [{thread_id}] 工作流执行失败: {e}")
        import traceback
        safe_print(traceback.format_exc())
        counter.increment_fail()
        return False

    # 6. 保存结果
    safe_print(f"\n[线程:{thread_id}][6] 保存结果...")
    try:
        # 使用帖子文件名作为前缀
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_filename = f"{post_name}_with_history_{timestamp}.json"
        output_path = output_dir / output_filename

        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(result, f, ensure_ascii=False, indent=2)

        safe_print(f"✅ [{thread_id}] 结果已保存到: {output_path}")
    except Exception as e:
        safe_print(f"❌ [{thread_id}] 保存结果失败: {e}")
        counter.increment_fail()
        return False

    # 7. 显示结果摘要
    safe_print(f"\n{'='*80}")
    safe_print(f"结果摘要 - {post_name} [线程:{thread_id}]")
    safe_print(f"{'='*80}")

    if result:
        three_points = result.get("三点解构", {})
        inspiration_data = three_points.get("灵感点", {})
        keypoints_data = three_points.get("关键点", {})
        comments = result.get("评论分析", {}).get("解构维度", [])

        safe_print(f"\n三点解构:")
        safe_print(f"   - 灵感点数量: {inspiration_data.get('total_count', 0)}")
        safe_print(f"   - 灵感点分析模式: {inspiration_data.get('analysis_mode', '未知')}")
        safe_print(f"   - 目的点数量: 1")
        safe_print(f"   - 关键点数量: {keypoints_data.get('total_count', 0)}")

        safe_print(f"\n评论分析:")
        safe_print(f"   - 解构维度数: {len(comments)}")

        topic_understanding = result.get("选题理解", {})
        if topic_understanding:
            topic_theme = topic_understanding.get("topic_theme", "")
            safe_print(f"\n选题理解:")
            safe_print(f"   - 选题主题: {topic_theme}")

    counter.increment_success()
    return True


def main():
    """主函数（并发版本，带历史帖子）"""
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='批量处理帖子的What解构工作流')
    parser.add_argument('directory', type=str, help='帖子目录名（如"阿里多多酱"或"G88818"）')
    args = parser.parse_args()

    directory = args.directory

    print("=" * 80)
    print(f"开始批量处理作者历史帖子（并发模式，带历史帖子）- 目录: {directory}")
    print("=" * 80)

    # 配置
    posts_dir = Path(__file__).parent / directory / "作者历史帖子"
    output_dir = Path(__file__).parent / directory / "output"
    output_dir.mkdir(parents=True, exist_ok=True)

    # 并发配置：设置最大线程数（建议根据CPU核心数和API限制调整）
    MAX_WORKERS = 4  # 可以根据需要调整，建议不超过5

    # 处理帖子数量限制
    MAX_POSTS = 20  # 只处理最新的20个帖子

    # 1. 加载所有帖子文件（按时间从新到旧排序，取最新20个）
    print(f"\n[1] 扫描帖子文件...")
    try:
        post_files = load_post_files(posts_dir, max_count=MAX_POSTS)
        print(f"✅ 选择 {len(post_files)} 个最新帖子进行处理")
    except Exception as e:
        print(f"❌ 扫描帖子文件失败: {e}")
        return

    # 2. 初始化线程安全计数器
    print(f"\n[2] 初始化并发处理...")
    print(f"   - 最大并发线程数: {MAX_WORKERS}")
    counter = ThreadSafeCounter()

    # 3. 使用线程池并发处理所有帖子
    print(f"\n[3] 开始并发处理...")
    print(f"   注意：多个线程会同时处理不同的帖子，每个帖子都会加载对应的历史帖子")
    print("=" * 80)

    start_time = datetime.now()

    # 使用ThreadPoolExecutor进行并发处理
    with ThreadPoolExecutor(max_workers=MAX_WORKERS, thread_name_prefix="Worker") as executor:
        # 提交所有任务
        future_to_post = {
            executor.submit(
                process_single_post,
                post_file,
                posts_dir,  # 传入帖子目录，用于加载历史帖子
                output_dir,
                counter,
                len(post_files),
                i
            ): (post_file, i)
            for i, post_file in enumerate(post_files, 1)
        }

        # 等待所有任务完成
        for future in as_completed(future_to_post):
            post_file, index = future_to_post[future]
            try:
                result = future.result()
                if not result:
                    safe_print(f"⚠️  处理帖子失败: {post_file.name}")
            except Exception as e:
                safe_print(f"❌ 处理帖子时发生异常: {post_file.name}")
                safe_print(f"   错误: {e}")
                import traceback
                safe_print(traceback.format_exc())
                counter.increment_fail()

    end_time = datetime.now()
    duration = (end_time - start_time).total_seconds()

    # 4. 总结
    print("\n" + "=" * 80)
    print("批量处理完成")
    print("=" * 80)
    print(f"\n总计: {len(post_files)} 个帖子")
    print(f"成功: {counter.success} 个")
    print(f"失败: {counter.fail} 个")
    print(f"耗时: {duration:.2f} 秒")
    print(f"平均每个帖子: {duration/len(post_files):.2f} 秒")
    print(f"\n结果保存目录: {output_dir}")
    print("=" * 80)


if __name__ == "__main__":
    main()