TravelContentCreator/scripts/distribution/extract_and_render.py

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

import os
import json
import shutil
import csv
import traceback
import re
import argparse
from datetime import datetime
import sqlite3
import logging

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

# 内置数据库记录功能
def init_database(db_path):
    """初始化数据库，创建表结构"""
    try:
        conn = sqlite3.connect(db_path)
        conn.execute("PRAGMA foreign_keys = OFF")  # 禁用外键约束
        cursor = conn.cursor()
        
        # 创建内容表
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS contents (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                entry_id TEXT NOT NULL UNIQUE,
                output_txt_path TEXT,
                poster_path TEXT,
                article_json_path TEXT,
                product TEXT,
                object TEXT,
                date TEXT,
                logic TEXT,
                judge_status INTEGER,
                is_distributed INTEGER DEFAULT 0,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)
        
        # 创建索引
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_contents_entry_id ON contents(entry_id)")
        
        conn.commit()
        logger.info("数据库初始化成功")
        return conn
    except sqlite3.Error as e:
        logger.error(f"初始化数据库失败: {e}")
        return None

def record_to_database(
    db_path, 
    entry_id, 
    output_txt_path=None, 
    poster_path=None, 
    article_json_path=None,
    product=None,
    object=None,
    date=None,
    logic=None,
    judge_status=None,
    is_distributed=0
):
    """将内容记录到数据库"""
    try:
        # 检查数据库是否存在，如果不存在则初始化
        if not os.path.exists(db_path):
            logger.info(f"数据库文件不存在: {db_path}，将自动创建")
            conn = init_database(db_path)
            if not conn:
                return False
        else:
            try:
                conn = sqlite3.connect(db_path)
                conn.execute("PRAGMA foreign_keys = OFF")  # 禁用外键约束
            except sqlite3.Error as e:
                logger.error(f"连接数据库失败: {e}")
                return False
        
        try:
            cursor = conn.cursor()
            
            # 准备数据
            data = (
                entry_id,
                output_txt_path or '',
                poster_path or '',
                article_json_path or '',
                product or '',
                object or '',
                date or '',
                logic or '',
                judge_status if judge_status is not None else None,
                is_distributed
            )
            
            # 插入或更新内容
            cursor.execute("""
                INSERT OR REPLACE INTO contents 
                (entry_id, output_txt_path, poster_path, article_json_path, 
                 product, object, date, logic, judge_status, is_distributed)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, data)
            
            conn.commit()
            logger.info(f"已将内容 {entry_id} 记录到数据库")
            return True
            
        except Exception as e:
            logger.error(f"记录内容到数据库失败: {e}")
            try:
                conn.rollback()
            except:
                pass
            return False
        finally:
            try:
                conn.close()
            except:
                pass
                
    except Exception as e:
        logger.error(f"记录提取内容时发生错误: {e}")
        return False

def convert_json_to_txt_content(json_path, prefer_original=False):
    """
    读取 JSON 文件，提取标题、内容和标签，移除 Markdown 格式，
    并返回格式化文本。
    
    根据judge_success字段决定使用原始内容还是审核后内容：
    - judge_success为True时使用title/content（除非prefer_original=True）
    - judge_success为False时使用original_title/original_content
    
    Args:
        json_path: JSON文件路径
        prefer_original: 是否优先使用原始内容，无视judge_success结果
    """
    print(f"    - 正在读取 JSON: {json_path}")
    if not os.path.exists(json_path):
        print(f"    - 警告: JSON 文件不存在: {json_path}")
        return None, f"文件未找到: {json_path}"

    try:
        with open(json_path, 'r', encoding='utf-8') as f_json:
            data = json.load(f_json)

        # 根据judge_success选择标题和内容
        judge_success = data.get('judge_success', None)
        
        if prefer_original and 'original_title' in data and 'original_content' in data:
            # 优先使用原始内容
            title = data.get('original_title', '未找到原始标题')
            content = data.get('original_content', '未找到原始内容')
            # 优先使用原始标签
            tags = data.get('original_tags', data.get('tags', '未找到标签'))
            print(f"    - 优先使用原始内容 (prefer_original=True)")
        elif judge_success is True and not prefer_original:
            # 使用审核后的内容
            title = data.get('title', '未找到标题')
            content = data.get('content', '未找到内容')
            tags = data.get('tags', '未找到标签')
            print(f"    - 使用审核后内容 (judge_success=True)")
        elif 'original_title' in data and 'original_content' in data:
            # 使用原始内容
            title = data.get('original_title', '未找到原始标题')
            content = data.get('original_content', '未找到原始内容')
            # 优先使用原始标签
            tags = data.get('original_tags', data.get('tags', '未找到标签'))
            print(f"    - 使用原始内容 (judge_success={judge_success})")
        else:
            # 若无original字段，使用常规字段
            title = data.get('title', '未找到标题')
            content = data.get('content', '未找到内容')
            tags = data.get('tags', '未找到标签')
            print(f"    - 使用常规内容 (无judge结果)")
        
        # 解决tag/tags字段重复问题，按照修正后的处理逻辑，只使用tags字段
        if not tags and 'tag' in data:
            tags = data.get('tag', '未找到标签')
            print(f"    - 使用tag字段作为标签 (该字段将在后续版本中统一为tags)")
        
        # 移除Markdown格式
        content_no_format = re.sub(r'\*\*(.*?)\*\*', r'\1', content)
        
        # 组合输出文本
        return f"{title}\n\n{content_no_format}\n\n{tags}", None
    except json.JSONDecodeError:
        print(f"    - 错误: JSON 格式无效: {json_path}")
        return None, f"无效的 JSON 格式: {json_path}"
    except Exception as e:
        print(f"    - 错误: 处理 JSON 时出错: {e}")
        return None, f"处理 JSON 时出错: {e}"

def load_topic_data(source_dir, run_id):
    """
    加载选题数据
    
    Args:
        source_dir: 源目录路径
        run_id: 运行ID
    
    Returns:
        dict: 以topic_index为键的选题数据字典
    """
    topic_file_path = os.path.join(source_dir, f"tweet_topic_{run_id}.json")
    topic_data = {}
    
    if os.path.exists(topic_file_path):
        try:
            with open(topic_file_path, 'r', encoding='utf-8') as f:
                topics = json.load(f)
                
            # 将选题数据转换为以index为键的字典
            for topic in topics:
                index = topic.get("index")
                if index:
                    topic_data[index] = topic
            
            print(f"成功加载选题数据，共{len(topic_data)}条")
        except Exception as e:
            print(f"加载选题数据时出错: {e}")
    else:
        print(f"警告: 未找到选题文件: {topic_file_path}")
    
    return topic_data

def process_result_directory(source_dir, output_dir, run_id=None, prefer_original=False, db_path=None):
    """
    处理指定的结果目录，提取内容并渲染到输出目录。
    
    Args:
        source_dir: 源目录路径，包含i_j子目录
        output_dir: 输出目录路径
        run_id: 可选的运行ID，如果不提供则使用源目录名
        prefer_original: 是否优先使用原始内容，无视judge_success结果
        db_path: 数据库路径，若不提供则使用默认路径
    """
    if not os.path.isdir(source_dir):
        print(f"错误: 源目录不存在: {source_dir}")
        return
    
    # 设置默认数据库路径
    if db_path is None:
        db_path = '/root/autodl-tmp/TravelContentCreator/distribution.db'
    
    # 数据库是否启用
    db_enabled = True
    
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    print(f"确保输出目录存在: {output_dir}")
    
    # 提取run_id
    if not run_id:
        run_id = os.path.basename(source_dir)
    
    # 加载选题数据
    topic_data = load_topic_data(source_dir, run_id)
    
    # 创建CSV清单，添加选题相关字段
    csv_path = os.path.join(output_dir, f"manifest_{run_id}.csv")
    csv_data = [
        [
            "EntryID",
            "TopicIndex",
            "VariantIndex",
            "Date",
            "Logic",
            "Object",
            "Product",
            "ProductLogic",
            "Style",
            "StyleLogic", 
            "TargetAudience",
            "TargetAudienceLogic",
            "SourcePath",
            "ArticleJsonPath",
            "OutputTxtPath",
            "PosterPath",
            "AdditionalImagesCount",
            "Status",
            "Details",
            "JudgeStatus",
            "ContentSource",
            "RecordedInDB",
            "IsDistributed"
        ]
    ]
    
    # 查找所有i_j目录
    entry_pattern = re.compile(r"^(\d+)_(\d+)$")
    entries = []
    
    for item in os.listdir(source_dir):
        item_path = os.path.join(source_dir, item)
        match = entry_pattern.match(item)
        if os.path.isdir(item_path) and match:
            entries.append(item)
    
    if not entries:
        print(f"警告: 在源目录中未找到任何i_j格式的子目录")
        return
    
    print(f"找到 {len(entries)} 个条目目录")
    
    # 处理每个条目
    for entry in sorted(entries):
        entry_path = os.path.join(source_dir, entry)
        output_entry_path = os.path.join(output_dir, entry)
        
        print(f"\n处理条目: {entry}")
        
        # 解析topic_index和variant_index
        match = entry_pattern.match(entry)
        topic_index = match.group(1)
        variant_index = match.group(2)
        
        # 获取该话题的选题信息
        topic_info = topic_data.get(topic_index, {})
        
        # 创建记录
        record = {
            "EntryID": entry,
            "TopicIndex": topic_index,
            "VariantIndex": variant_index,
            "Date": topic_info.get("date", ""),
            "Logic": topic_info.get("logic", ""),
            "Object": topic_info.get("object", ""),
            "Product": topic_info.get("product", ""),
            "ProductLogic": topic_info.get("product_logic", ""),
            "Style": topic_info.get("style", ""),
            "StyleLogic": topic_info.get("style_logic", ""),
            "TargetAudience": topic_info.get("target_audience", ""),
            "TargetAudienceLogic": topic_info.get("target_audience_logic", ""),
            "SourcePath": entry_path,
            "ArticleJsonPath": "",
            "OutputTxtPath": "",
            "PosterPath": "",
            "AdditionalImagesCount": 0,
            "Status": "Processing",
            "Details": "",
            "JudgeStatus": "",
            "ContentSource": "unknown",
            "RecordedInDB": "No",
            "IsDistributed": "No"
        }
        
        # 创建输出条目目录
        try:
            os.makedirs(output_entry_path, exist_ok=True)
        except Exception as e:
            record["Status"] = "Failed"
            record["Details"] = f"创建输出目录失败: {e}"
            csv_data.append([record[col] for col in csv_data[0]])
            print(f"  - 错误: {record['Details']}")
            continue
        
        # 1. 处理article.json -> txt
        json_path = os.path.join(entry_path, "article.json")
        txt_path = os.path.join(output_entry_path, "article.txt")
        record["ArticleJsonPath"] = json_path
        record["OutputTxtPath"] = txt_path
        
        # 读取article.json
        article_data = {}
        if os.path.exists(json_path):
            try:
                with open(json_path, 'r', encoding='utf-8') as f_json:
                    article_data = json.load(f_json)
                    # 提取judge_success状态
                    if "judge_success" in article_data:
                        record["JudgeStatus"] = str(article_data["judge_success"])
                    elif "judged" in article_data:
                        record["JudgeStatus"] = "已审核" if article_data["judged"] else "未审核"
            except Exception as e:
                print(f"  - 错误: 读取article.json失败: {e}")
                
            txt_content, error = convert_json_to_txt_content(json_path, prefer_original)
            if error:
                record["Status"] = "Partial"
                record["Details"] += f"文章处理失败: {error}; "
                print(f"  - 错误: {record['Details']}")
            else:
                try:
                    with open(txt_path, 'w', encoding='utf-8') as f_txt:
                        f_txt.write(txt_content)
                    print(f"  - 成功写入文本文件: {txt_path}")
                    
                    # 记录内容来源
                    if prefer_original:
                        record["ContentSource"] = "original_preferred"
                    elif article_data.get("judge_success") is True:
                        record["ContentSource"] = "judged"
                    elif "original_title" in article_data:
                        record["ContentSource"] = "original"
                    else:
                        record["ContentSource"] = "default"
                        
                except Exception as e:
                    record["Status"] = "Partial"
                    record["Details"] += f"写入文本文件失败: {e}; "
                    print(f"  - 错误: {record['Details']}")
        else:
            record["Status"] = "Partial"
            record["Details"] += "文章JSON文件不存在; "
            print(f"  - 警告: {record['Details']}")
        
        # 2. 处理海报图片
        poster_dir = os.path.join(entry_path, "poster")
        poster_jpg_path = os.path.join(poster_dir, "poster.jpg")
        output_poster_path = os.path.join(output_entry_path, "poster.jpg")
        record["PosterPath"] = output_poster_path
        
        if os.path.exists(poster_jpg_path):
            try:
                shutil.copy2(poster_jpg_path, output_poster_path)
                print(f"  - 成功复制海报图片: {output_poster_path}")
            except Exception as e:
                record["Status"] = "Partial"
                record["Details"] += f"复制海报图片失败: {e}; "
                print(f"  - 错误: {record['Details']}")
        else:
            record["Status"] = "Partial"
            record["Details"] += "海报图片不存在; "
            print(f"  - 警告: {record['Details']}")
        
        # 3. 处理额外图片
        image_dir = os.path.join(entry_path, "image")
        output_image_dir = os.path.join(output_entry_path, "additional_images")
        
        if os.path.exists(image_dir) and os.path.isdir(image_dir):
            try:
                os.makedirs(output_image_dir, exist_ok=True)
                image_count = 0
                
                for filename in os.listdir(image_dir):
                    if filename.startswith("additional_") and filename.endswith(".jpg"):
                        source_file = os.path.join(image_dir, filename)
                        dest_file = os.path.join(output_image_dir, filename)
                        
                        # 复制图片
                        shutil.copy2(source_file, dest_file)
                        image_count += 1
                
                record["AdditionalImagesCount"] = image_count
                print(f"  - 复制了 {image_count} 张额外图片到: {output_image_dir}")
            except Exception as e:
                record["Status"] = "Partial"
                record["Details"] += f"处理额外图片时出错: {e}; "
                print(f"  - 错误: {record['Details']}")
        else:
            record["AdditionalImagesCount"] = 0
            print(f"  - 没有找到额外图片目录")
        
        # 更新状态
        if record["Status"] == "Processing":
            record["Status"] = "Success"
            record["Details"] = "处理成功完成"
        
        # 4. 将内容记录到数据库
        if db_enabled:
            try:
                # 准备judge_status值
                if record["JudgeStatus"] == "True":
                    judge_status = 1
                elif record["JudgeStatus"] == "False":
                    judge_status = 0
                else:
                    judge_status = None
                
                # 调用数据库记录函数
                success = record_to_database(
                    db_path,
                    entry_id=record["EntryID"],
                    output_txt_path=record["OutputTxtPath"],
                    poster_path=record["PosterPath"],
                    article_json_path=record["ArticleJsonPath"],
                    product=record["Product"],
                    object=record["Object"],
                    date=record["Date"],
                    logic=record["Logic"],
                    judge_status=judge_status,
                    is_distributed=0  # 默认为未分发
                )
                
                if success:
                    record["RecordedInDB"] = "Yes"
                    print(f"  - 成功将内容记录到数据库")
                else:
                    record["RecordedInDB"] = "Failed"
                    print(f"  - 警告: 内容记录到数据库失败")
            except Exception as e:
                record["RecordedInDB"] = "Error"
                print(f"  - 错误: 记录到数据库时发生异常: {e}")
                traceback.print_exc()  # 打印详细的异常堆栈
        else:
            record["RecordedInDB"] = "Disabled"
            print(f"  - 信息: 数据库记录功能已禁用")
        
        # 添加记录到CSV数据
        csv_data.append([record[col] for col in csv_data[0]])
    
    # 写入CSV清单
    try:
        print(f"\n正在写入清单CSV: {csv_path}")
        with open(csv_path, 'w', newline='', encoding='utf-8-sig') as f_csv:
            writer = csv.writer(f_csv)
            writer.writerows(csv_data)
        print(f"清单CSV生成成功")
    except Exception as e:
        print(f"写入CSV文件时出错: {e}")
        traceback.print_exc()
    
    print(f"\n处理完成. 共处理 {len(entries)} 个条目.")
    print(f"结果保存在: {output_dir}")

def main():
    parser = argparse.ArgumentParser(description="从TravelContentCreator结果目录提取内容并渲染到指定目录")
    parser.add_argument("--source", type=str, help="源目录路径")
    parser.add_argument("--output", type=str, help="输出目录路径")
    parser.add_argument("--run-id", type=str, help="自定义运行ID")
    parser.add_argument("--prefer-original", action="store_true", help="优先使用原始内容，忽略审核结果")
    parser.add_argument("--db-path", type=str, help="数据库路径，若不提供则使用默认路径")
    parser.add_argument("--disable-db", action="store_true", help="禁用数据库记录功能")
    
    args = parser.parse_args()
    
    # 默认值设置
    source = args.source if args.source else "/root/autodl-tmp/TravelContentCreator/result/2025-05-14_22-10-37"
    output = args.output if args.output else "/root/autodl-tmp/TravelContentCreator/output/2025-05-14_22-10-37"
    run_id = args.run_id if args.run_id else os.path.basename(source)
    prefer_original = args.prefer_original
    db_path = args.db_path if args.db_path else '/root/autodl-tmp/TravelContentCreator/distribution.db'
    
    print("-" * 60)
    print(f"开始提取和渲染流程")
    print(f"源目录: {source}")
    print(f"输出目录: {output}")
    print(f"运行ID: {run_id}")
    if prefer_original:
        print("内容模式: 优先使用原始内容")
    else:
        print("内容模式: 根据审核结果选择内容")
    
    if args.disable_db:
        print("数据库记录: 已禁用")
    else:
        print(f"数据库记录: 已启用 (路径: {db_path})")
    print("-" * 60)
    
    process_result_directory(source, output, run_id, prefer_original, db_path)
    
    print("\n脚本执行完毕.")

if __name__ == "__main__":
    main()