main.py

#!/usr/bin/env python3
"""
公众号爆款生成器 - 主程序
从选题到成稿一键完成
支持自我迭代优化
"""

import os
import sys
import json
import argparse
from pathlib import Path
from datetime import datetime

# 加载环境变量
def load_env_file(env_file: Path):
    """加载 .env 文件中的环境变量"""
    if env_file.exists():
        with open(env_file, 'r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith('#') and '=' in line:
                    key, value = line.split('=', 1)
                    os.environ[key.strip()] = value.strip()

# 加载 Get 笔记配置
config_dir = Path(__file__).parent / "config"
getnote_env = config_dir / "getnote.env"
load_env_file(getnote_env)

# 导入各模块
from mcp_tools.feishu import FeishuAPI
from mcp_tools.getnote import GetNoteAPI
from modules.image_match import ImageMatcher
from modules.layout_engine import WeChatLayoutEngine
from modules.content_gen import ContentGenerator


class OptimizedContentGenerator(ContentGenerator):
    """优化版内容生成器 - 应用自我学习的配置"""

    def __init__(self):
        super().__init__()
        self.config_dir = Path(__file__).parent / "config"
        self.prompt_templates = self.config_dir / "prompt_templates.json"
        self._load_optimized_config()

    def _load_optimized_config(self):
        """加载优化后的配置"""
        if self.prompt_templates.exists():
            try:
                with open(self.prompt_templates, 'r', encoding='utf-8') as f:
                    self.templates = json.load(f)
                print(f"   ✅ 已加载优化配置 v{self.templates['content_generation']['version']}")
            except Exception as e:
                print(f"   ⚠️  加载优化配置失败: {e}")
                self.templates = None
        else:
            self.templates = None

    def _build_prompt(self, topic: str, category: str, materials: list,
                     style: str, word_count: int) -> str:
        """构建优化后的提示词"""

        # 如果有优化配置，使用优化后的提示词
        if self.templates and 'content_generation' in self.templates:
            config = self.templates['content_generation']
            persona = config.get('persona', {})
            writing_rules = config.get('writing_rules', [])
            avoid_list = config.get('avoid', [])

            # 构建人设部分
            persona_text = f"""你是 {persona.get('name', 'ruby鑫燕')}，一位高级经管硕士背景的职场专家，同时也是个地道的北京大妞儿。

**你的人设特点**：
{chr(10).join(f'- {trait}' for trait in persona.get('traits', []))}

**语言风格**：
{chr(10).join(f'- {style_item}' for style_item in persona.get('language_style', []))}

**写作要求**：
{chr(10).join(f'{i+1}. {rule}' for i, rule in enumerate(writing_rules))}

**务必避免**：
{chr(10).join(f'❌ {avoid}' for avoid in avoid_list)}
"""
        else:
            # 使用默认提示词
            persona_text = f"""你是一个专业的公众号爆款文案写手。

**要求**:
1. 标题要吸引人，能激发读者的好奇心或共鸣
2. 开头要有钩子，快速抓住读者注意力
3. 内容要有价值，提供实用的观点或方法
4. 结构清晰，使用小标题划分段落
5. 语言简洁生动，符合微信公众号的阅读习惯
6. 结尾要有行动号召或情感共鸣
"""

        prompt = f"""{persona_text}

**主题**: {topic}
**分类**: {category}
**风格**: {style}
**字数**: 约{word_count}字
"""

        # 如果有参考素材，加入提示词
        if materials and len(materials) > 0:
            prompt += "\n**参考素材**:\n"
            for idx, material in enumerate(materials[:3], 1):  # 最多3个素材
                prompt += f"\n{idx}. {material.get('title', '')}\n"
                # 只取素材的部分内容
                content_snippet = material.get('content', '')[:200]
                prompt += f"   {content_snippet}...\n"

        prompt += """

请以 Markdown 格式输出，包含：
- 标题（h1）
- 3-5个小节（h2）
- 每个小节下有2-3个要点

开始创作："""

        return prompt


class WeChatBurstGen:
    """公众号爆款内容生成器"""

    def __init__(
        self,
        feishu_app_id: str = None,
        feishu_app_secret: str = None,
        getnote_api_key: str = None,
        getnote_kb_id: str = None,
        style: str = "简约",
        skip_feishu: bool = False,
        skip_getnote: bool = False
    ):
        """初始化生成器"""
        print("\n🚀 初始化公众号爆款生成器...\n")

        # 初始化各模块
        self.feishu = FeishuAPI(feishu_app_id, feishu_app_secret, skip_validation=skip_feishu)

        # 初始化 Get 笔记
        self.getnote = None
        if not skip_getnote:
            try:
                self.getnote = GetNoteAPI(api_key=getnote_api_key, kb_id=getnote_kb_id)
                print("   ✅ Get 笔记 API 已连接")
            except Exception as e:
                print(f"   ⚠️  Get 笔记初始化失败: {e}")
                print("   💡 将跳过素材获取，仅使用 AI 生成")

        self.image_matcher = ImageMatcher()
        self.layout_engine = WeChatLayoutEngine(style=style)
        self.content_gen = OptimizedContentGenerator()  # 使用优化版生成器

        self.style = style
        self.output_dir = "/Users/rubyliu/Desktop/wechatknow"
        Path(self.output_dir).mkdir(parents=True, exist_ok=True)

        print("   ✅ 飞书 API 已连接")
        print("   ✅ 图片生成器已就绪")
        print("   ✅ 排版引擎已加载")
        print("   ✅ AI 文案生成器已就绪\n")

    def _sanitize_filename(self, title: str, max_length: int = 30) -> str:
        """
        将标题转换为安全的文件名
        - 保留中文、英文、数字
        - 移除特殊字符
        - 限制长度
        """
        import re

        # 移除标点符号和特殊字符，保留中英文、数字
        safe_title = re.sub(r'[^\w\s\u4e00-\u9fff]', '', title)
        # 移除多余空格
        safe_title = re.sub(r'\s+', '_', safe_title.strip())
        # 限制长度
        if len(safe_title) > max_length:
            safe_title = safe_title[:max_length]

        return safe_title if safe_title else "untitled"

    def _generate_markdown(self, title: str, author: str, content: str,
                          cover_image: str = None, inline_images: list = None) -> str:
        """生成 Markdown 格式"""
        # 调试信息
        print(f"   🔧 _generate_markdown 调试:")
        print(f"      - cover_image: {cover_image}")
        print(f"      - inline_images: {inline_images}")

        md = f"# {title}\n\n"
        md += f"**作者：{author}**\n\n---\n\n"

        if cover_image:
            md += f"![封面图]({cover_image})\n\n"
            print(f"      - ✅ 已添加封面图: {cover_image}")
        else:
            print(f"      - ⚠️  未添加封面图 (cover_image 为 None)")

        # 处理内容，插入配图
        lines = content.split('\n')
        result_lines = []
        image_index = 0
        inline_images = inline_images or []

        if inline_images:
            print(f"      - 准备插入 {len(inline_images)} 张配图")

        for line in lines:
            result_lines.append(line)

            # 在二级标题后插入图片
            if line.startswith('## ') and image_index < len(inline_images):
                img_path = inline_images[image_index]
                result_lines.append(f"\n![配图{image_index + 1}]({img_path})\n")
                print(f"      - ✅ 在标题 '{line[:30]}...' 后插入配图{image_index + 1}: {img_path}")
                image_index += 1

        if inline_images and image_index < len(inline_images):
            print(f"      - ⚠️  警告: 有 {len(inline_images) - image_index} 张配图未能插入 (文章二级标题不足)")

        md += '\n'.join(result_lines)
        md += "\n\n---\n\n**— END —**\n"

        return md

    def _generate_xiumi(self, title: str, author: str, content: str,
                       cover_image: str = None, inline_images: list = None) -> str:
        """生成秀米编辑器格式"""
        xiumi = '<section style="margin: 10px auto; font-size: 16px; color: #333; line-height: 1.8; max-width: 100%;">\n\n'

        # 封面图
        if cover_image:
            xiumi += '<section style="text-align: center; margin-bottom: 20px;">\n'
            xiumi += '<img src="图片1" style="width: 100%; max-width: 100%; height: auto; display: block;" />\n'
            xiumi += '</section>\n\n'

        # 标题
        xiumi += f'<h1 style="font-size: 24px; font-weight: bold; text-align: center; margin: 20px 0; color: #1a1a1a;">{title}</h1>\n\n'

        # 作者
        xiumi += f'<p style="text-align: center; color: #999; font-size: 14px; margin-bottom: 30px;">作者：{author}</p>\n\n'

        # 处理内容
        lines = content.split('\n')
        image_index = 1  # 封面已用图片1
        inline_images = inline_images or []

        for line in lines:
            line = line.strip()
            if not line:
                continue

            if line.startswith('# '):
                continue  # 跳过一级标题
            elif line.startswith('## '):
                # 二级标题
                h2_text = line.replace('## ', '')
                xiumi += f'<h2 style="font-size: 20px; font-weight: bold; margin: 30px 0 15px 0; border-left: 4px solid #4CAF50; padding-left: 10px;">{h2_text}</h2>\n\n'

                # 插入配图
                if image_index < len(inline_images) + 1:
                    image_index += 1
                    xiumi += f'<section style="margin: 20px 0; text-align: center;">\n'
                    xiumi += f'<img src="图片{image_index}" style="width: 100%; max-width: 100%; height: auto; display: block;" />\n'
                    xiumi += '</section>\n\n'

            elif line.startswith('### '):
                # 三级标题
                h3_text = line.replace('### ', '')
                xiumi += f'<h3 style="font-size: 18px; font-weight: bold; margin: 25px 0 12px 0; color: #2c3e50;">{h3_text}</h3>\n\n'

            elif line.startswith('```'):
                # 代码块开始/结束
                continue

            elif line.startswith('- '):
                # 列表项
                li_text = line.replace('- ', '')
                xiumi += f'<li style="margin: 8px 0;">{li_text}</li>\n'

            elif line.startswith(('1. ', '2. ', '3. ', '4. ', '5. ', '6. ', '7. ', '8. ', '9. ')):
                # 有序列表
                li_text = line[3:]
                xiumi += f'<li style="margin: 8px 0;">{li_text}</li>\n'

            else:
                # 普通段落
                if '**' in line:
                    line = line.replace('**', '<strong>').replace('**', '</strong>')
                xiumi += f'<p style="text-indent: 2em; margin: 15px 0;">{line}</p>\n\n'

        xiumi += '<p style="text-align: center; color: #999; font-size: 12px; margin-top: 30px;">— END —</p>\n\n'
        xiumi += '</section>'

        return xiumi

    def generate_complete_article(
        self,
        topic: str,
        category: str = "",
        word_count: int = 1500,
        save_to_feishu: bool = True,
        auto_review: bool = False,
        auto_optimize: bool = False
    ) -> dict:
        """
        完整的文章生成流程

        Args:
            topic: 选题主题
            category: 文章分类
            word_count: 目标字数
            save_to_feishu: 是否保存到飞书
            auto_review: 是否自动审稿
            auto_optimize: 是否自动优化配置

        Returns:
            包含文章信息和文件路径的字典
        """
        print("=" * 70)
        print(f"📝 开始生成文章")
        print(f"   选题: {topic}")
        print(f"   分类: {category}")
        print(f"   风格: {self.style}")
        print("=" * 70)

        # 步骤 1: 获取参考素材（从 Get 笔记）
        materials = []
        if self.getnote:
            print("\n【步骤 1/6】从 Get 笔记获取参考素材")
            try:
                materials = self.getnote.get_materials_for_topic(
                    topic=topic,
                    category=category,
                    limit=5
                )
                if materials:
                    print(f"   ✅ 获取到 {len(materials)} 条参考素材")
                else:
                    print("   ⚠️  未找到相关素材，将纯 AI 生成")
            except Exception as e:
                print(f"   ⚠️  获取素材失败: {e}")
                print("   💡 将继续使用纯 AI 生成")
        else:
            print("\n【步骤 1/6】跳过素材获取（未配置 Get 笔记）")

        # 步骤 2: AI 生成文案
        print("\n【步骤 2/6】生成文案内容")
        article = self.content_gen.generate_article(
            topic=topic,
            category=category,
            materials=materials,
            style=self.style,
            word_count=word_count
        )

        title = article['title']
        content = article['content']
        sections = article['sections']

        # 步骤 3: 生成封面图
        print("\n【步骤 3/6】生成封面图片")
        cover_image = self.image_matcher.generate_cover(
            title=title,
            category=category,
            style=self.style
        )

        if cover_image:
            print(f"   ✅ 封面图: {cover_image}")
        else:
            print(f"   ⚠️  未生成封面图")

        # 步骤 4: 生成配图
        print("\n【步骤 4/6】生成文章配图")
        inline_images = self.image_matcher.generate_inline_images(
            sections=sections,
            count=min(3, len(sections)),
            style=self.style
        )

        if inline_images:
            print(f"   ✅ 生成配图: {len(inline_images)} 张")
        else:
            print(f"   ⚠️  未生成配图")

        # 步骤 5: 生成多格式输出
        print("\n【步骤 5/6】生成多格式文件")
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

        # 创建 images 子目录
        images_dir = os.path.join(self.output_dir, "images")
        Path(images_dir).mkdir(parents=True, exist_ok=True)

        # 复制图片到 images 子目录
        import shutil
        copied_cover = None
        copied_inline_images = []

        # 检查并复制封面图
        if cover_image:
            print(f"   🔍 检查封面图: {cover_image}")
            if not os.path.exists(cover_image):
                print(f"   ❌ 错误: 封面图文件不存在: {cover_image}")
            else:
                cover_filename = os.path.basename(cover_image)
                dest_cover = os.path.join(images_dir, cover_filename)
                try:
                    shutil.copy2(cover_image, dest_cover)
                    copied_cover = f"images/{cover_filename}"  # 使用相对路径
                    print(f"   ✅ 封面图已复制: {cover_filename} -> {copied_cover}")
                except Exception as e:
                    print(f"   ❌ 复制封面图失败: {e}")
                    copied_cover = None
        else:
            print(f"   ⚠️  警告: 没有生成封面图")

        # 检查并复制配图
        if inline_images:
            print(f"   🔍 检查配图: 共 {len(inline_images)} 张")
            for i, img_path in enumerate(inline_images, 1):
                if not os.path.exists(img_path):
                    print(f"   ❌ 错误: 配图{i}文件不存在: {img_path}")
                    continue

                img_filename = os.path.basename(img_path)
                dest_img = os.path.join(images_dir, img_filename)
                try:
                    shutil.copy2(img_path, dest_img)
                    relative_path = f"images/{img_filename}"
                    copied_inline_images.append(relative_path)
                    print(f"   ✅ 配图{i}已复制: {img_filename} -> {relative_path}")
                except Exception as e:
                    print(f"   ❌ 复制配图{i}失败: {e}")

            if not copied_inline_images:
                print(f"   ⚠️  警告: 所有配图复制都失败了")
            else:
                print(f"   ✅ 成功复制 {len(copied_inline_images)}/{len(inline_images)} 张配图")
        else:
            print(f"   ⚠️  警告: 没有生成配图")

        # 5.1 保存 Markdown 文件
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        safe_title = self._sanitize_filename(title)
        md_filename = f"{safe_title}_{timestamp}.md"
        md_filepath = os.path.join(self.output_dir, md_filename)
        print(f"   📝 文件名: {md_filename}")

        md_content = self._generate_markdown(
            title=title,
            author="ruby鑫燕",
            content=content,
            cover_image=copied_cover,
            inline_images=copied_inline_images
        )

        with open(md_filepath, 'w', encoding='utf-8') as f:
            f.write(md_content)
        print(f"   ✅ Markdown 文件: {md_filepath}")

        # 验证图片路径是否正确插入到 Markdown 中
        image_count_in_md = md_content.count('![')
        expected_count = (1 if copied_cover else 0) + len(copied_inline_images)
        if image_count_in_md != expected_count:
            print(f"   ⚠️  警告: Markdown中的图片数量({image_count_in_md})与预期({expected_count})不符!")
            if copied_cover:
                if copied_cover not in md_content:
                    print(f"   ❌ 封面图路径未插入: {copied_cover}")
            for i, img_path in enumerate(copied_inline_images, 1):
                if img_path not in md_content:
                    print(f"   ❌ 配图{i}路径未插入: {img_path}")
        else:
            print(f"   ✅ 图片路径验证通过: {image_count_in_md} 张图片已正确插入")

        # 5.2 生成 HTML 排版
        html_content = self.layout_engine.render(
            title=title,
            author="ruby鑫燕",
            content=content,
            cover_image=cover_image,
            inline_images=inline_images
        )

        html_filename = f"{safe_title}_{timestamp}.html"
        html_filepath = os.path.join(self.output_dir, html_filename)

        with open(html_filepath, 'w', encoding='utf-8') as f:
            f.write(html_content)
        print(f"   ✅ HTML 文件: {html_filepath}")

        # 5.3 生成秀米格式
        xiumi_filename = f"{safe_title}_{timestamp}_xiumi.html"
        xiumi_filepath = os.path.join(self.output_dir, xiumi_filename)

        xiumi_content = self._generate_xiumi(
            title=title,
            author="ruby鑫燕",
            content=content,
            cover_image=cover_image,
            inline_images=inline_images
        )

        with open(xiumi_filepath, 'w', encoding='utf-8') as f:
            f.write(xiumi_content)
        print(f"   ✅ 秀米格式: {xiumi_filepath}")

        # 步骤 6: 保存到飞书
        doc_url = None
        if save_to_feishu:
            print("\n【步骤 6/6】保存到飞书")
            try:
                # 创建飞书文档
                doc_url = self.feishu.save_draft(title, html_content)
                print(f"   ✅ 飞书文档: {doc_url}")

                # 添加选题记录到多维表格
                try:
                    self.feishu.add_topic(
                        topic=topic,
                        category=category,
                        analysis=f"已生成文章，字数: {len(content)}",
                        status="已生成"
                    )
                    print(f"   ✅ 选题已记录到多维表格")
                except Exception as e:
                    print(f"   ⚠️  记录选题失败（可忽略）: {e}")

            except Exception as e:
                print(f"   ❌ 保存到飞书失败: {e}")
                print(f"   💡 提示: 请检查飞书配置或使用本地HTML文件")

        # 返回结果
        result = {
            'title': title,
            'content': content,
            'sections': sections,
            'word_count': len(content),
            'cover_image': cover_image,
            'inline_images': inline_images,
            'md_file': md_filepath,
            'html_file': html_filepath,
            'xiumi_file': xiumi_filepath,
            'feishu_doc': doc_url,
            'timestamp': timestamp
        }

        print("\n" + "=" * 70)
        print("✅ 文章生成完成！")
        print("=" * 70)
        print(f"📌 标题: {title}")
        print(f"📊 字数: {len(content)}")
        print(f"🖼️  封面: {cover_image or '无'}")
        print(f"🎨 配图: {len(inline_images)} 张")
        print(f"📄 Markdown: {md_filepath}")
        print(f"📁 HTML: {html_filepath}")
        print(f"✨ 秀米: {xiumi_filepath}")
        if doc_url:
            print(f"☁️  飞书: {doc_url}")
        print("=" * 70 + "\n")

        # 自动审稿
        review_score = None
        review_file = None
        if auto_review:
            print("\n" + "=" * 70)
            print("🔍 自动审稿")
            print("=" * 70 + "\n")

            try:
                from editor_review import SeniorEditor
                editor = SeniorEditor()
                review_result = editor.review_article(html_filepath)
                review_file = review_result.get('review_file')

                # 提取分数
                if review_file and os.path.exists(review_file):
                    with open(review_file, 'r', encoding='utf-8') as f:
                        review_content = f.read()
                        import re
                        score_match = re.search(r'总分.*?(\d+)/100', review_content)
                        if score_match:
                            review_score = int(score_match.group(1))
                            print(f"\n📊 审稿评分: {review_score}/100\n")

                            if auto_optimize and review_score < 85:
                                print("💡 分数低于85，启动自我优化...\n")

                                try:
                                    from self_learning import SelfLearningSystem
                                    learning_system = SelfLearningSystem()
                                    learning_system.analyze_review_and_optimize(review_file)
                                    print("\n✅ 配置已优化，下次生成将应用新配置\n")
                                except Exception as e:
                                    print(f"⚠️  自我优化失败: {e}\n")

            except Exception as e:
                print(f"⚠️  自动审稿失败: {e}\n")

        result = {
            'title': title,
            'content': content,
            'sections': sections,
            'word_count': len(content),
            'cover_image': cover_image,
            'inline_images': inline_images,
            'md_file': md_filepath,
            'html_file': html_filepath,
            'xiumi_file': xiumi_filepath,
            'feishu_doc': doc_url,
            'timestamp': timestamp,
            'review_score': review_score,
            'review_file': review_file
        }

        return result

    def batch_generate(
        self,
        topics: list,
        category: str = "",
        word_count: int = 1500
    ) -> list:
        """
        批量生成文章

        Args:
            topics: 选题列表
            category: 统一分类
            word_count: 目标字数

        Returns:
            结果列表
        """
        results = []

        print(f"\n📚 批量生成模式")
        print(f"   选题数量: {len(topics)}")
        print(f"   分类: {category}\n")

        for idx, topic in enumerate(topics, 1):
            print(f"\n{'='*70}")
            print(f"进度: {idx}/{len(topics)}")
            print(f"{'='*70}\n")

            try:
                result = self.generate_complete_article(
                    topic=topic,
                    category=category,
                    word_count=word_count
                )
                results.append(result)

                # 批量模式下稍作延时，避免API限流
                if idx < len(topics):
                    import time
                    print("\n⏳ 等待 5 秒后继续...\n")
                    time.sleep(5)

            except Exception as e:
                print(f"\n❌ 生成失败: {e}\n")
                results.append({
                    'topic': topic,
                    'error': str(e)
                })

        print(f"\n{'='*70}")
        print(f"✅ 批量生成完成！")
        print(f"   成功: {len([r for r in results if 'error' not in r])}/{len(topics)}")
        print(f"{'='*70}\n")

        return results


# ============ 命令行接口 ============

def main():
    parser = argparse.ArgumentParser(
        description='公众号爆款生成器 - 从选题到成稿一键完成'
    )

    parser.add_argument('topic', nargs='?', help='文章主题')
    parser.add_argument('-c', '--category', default='', help='文章分类（职场/生活/科技等）')
    parser.add_argument('-s', '--style', default='简约',
                       choices=['简约', '商务', '温馨', '科技'],
                       help='排版风格')
    parser.add_argument('-w', '--word-count', type=int, default=1500,
                       help='目标字数')
    parser.add_argument('--no-feishu', action='store_true',
                       help='不保存到飞书（仅生成本地文件）')
    parser.add_argument('--auto-review', action='store_true',
                       help='自动审稿')
    parser.add_argument('--auto-optimize', action='store_true',
                       help='自动优化配置（需配合 --auto-review）')

    # 批量模式
    parser.add_argument('--batch', help='批量生成，从文件读取选题（每行一个）')

    # 飞书配置
    parser.add_argument('--feishu-app-id', help='飞书 App ID')
    parser.add_argument('--feishu-app-secret', help='飞书 App Secret')

    args = parser.parse_args()

    # 初始化生成器
    try:
        generator = WeChatBurstGen(
            feishu_app_id=args.feishu_app_id,
            feishu_app_secret=args.feishu_app_secret,
            style=args.style,
            skip_feishu=args.no_feishu
        )
    except Exception as e:
        print(f"❌ 初始化失败: {e}")
        print("\n💡 提示: 请配置飞书环境变量或使用 --no-feishu 模式")
        print("   详见: config/FEISHU_SETUP.md\n")
        sys.exit(1)

    # 批量模式
    if args.batch:
        if not os.path.exists(args.batch):
            print(f"❌ 文件不存在: {args.batch}")
            sys.exit(1)

        with open(args.batch, 'r', encoding='utf-8') as f:
            topics = [line.strip() for line in f if line.strip()]

        generator.batch_generate(
            topics=topics,
            category=args.category,
            word_count=args.word_count
        )

    # 单篇模式
    elif args.topic:
        generator.generate_complete_article(
            topic=args.topic,
            category=args.category,
            word_count=args.word_count,
            save_to_feishu=not args.no_feishu,
            auto_review=args.auto_review,
            auto_optimize=args.auto_optimize
        )

    else:
        parser.print_help()
        print("\n示例:")
        print('  python main.py "如何提升工作效率" -c 职场 -s 简约')
        print('  python main.py "时间管理技巧" -c 职场 --auto-review --auto-optimize')
        print('  python main.py --batch topics.txt -c 生活\n')


if __name__ == "__main__":
    main()