#!/usr/bin/env python3
"""
知乎文章爬取脚本 - 基于 MediaCrawler
用于爬取知乎专栏文章并导出到知识库格式
"""
import argparse
import json
import os
import re
import sys
import subprocess
import shutil
from pathlib import Path
from datetime import datetime

# 配置
MEDIACRAWLER_DIR = Path.home() / "work" / "MediaCrawler"
DEFAULT_OUTPUT_DIR = Path.home() / "work" / "knowledge-base" / "data" / "zhihu"


def load_env():
    """加载环境变量"""
    env_file = Path.home() / ".kb_env"
    if env_file.exists():
        with open(env_file) as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    key, val = line.split("=", 1)
                    os.environ[key.strip()] = val.strip()


def get_zhihu_cookie():
    """获取知乎 Cookie"""
    return os.environ.get("ZHIHU_COOKIE", "")


def crawl_zhihu_article(article_url: str, output_dir: Path) -> dict:
    """
    使用 MediaCrawler 爬取知乎文章
    
    Args:
        article_url: 知乎文章链接
        output_dir: 输出目录
    
    Returns:
        dict: 爬取结果
    """
    # 清理旧数据
    data_dir = MEDIACRAWLER_DIR / "data" / "zhihu" / "json"
    data_dir.mkdir(parents=True, exist_ok=True)
    
    # 备份现有数据
    backup_dir = data_dir / "backup"
    backup_dir.mkdir(exist_ok=True)
    for f in data_dir.glob("detail_contents_*.json"):
        shutil.move(str(f), str(backup_dir / f.name))
    
    # 配置 MediaCrawler
    config_file = MEDIACRAWLER_DIR / "config" / "zhihu_config.py"
    
    # 读取配置文件并修改
    with open(config_file) as f:
        config_content = f.read()
    
    # 强制更新配置中的 ID 列表（每次都更新，确保爬取正确的文章）
    article_id = article_url.split("/")[-1]
    
    # 使用正则表达式替换整个 ZHIHU_SPECIFIED_ID_LIST 部分
    pattern = r'(ZHIHU_SPECIFIED_ID_LIST\s*=\s*\[)([\s\S]*?)(\])'
    replacement = f'\\1\n    "{article_url}",  # article\n\\3'
    
    new_config = re.sub(pattern, replacement, config_content)
    
    # 如果正则替换失败（没有匹配到），使用简单替换
    if new_config == config_content:
        # 直接替换整个配置项
        if 'ZHIHU_SPECIFIED_ID_LIST' in config_content:
            # 找到配置项的范围并替换
            lines = config_content.split('\n')
            new_lines = []
            in_list = False
            for line in lines:
                if 'ZHIHU_SPECIFIED_ID_LIST' in line and '=' in line:
                    new_lines.append(f'ZHIHU_SPECIFIED_ID_LIST = ["{article_url}",  # article')
                    in_list = True
                elif in_list and line.strip().startswith(']'):
                    new_lines.append(line)
                    in_list = False
                elif not in_list:
                    new_lines.append(line)
            new_config = '\n'.join(new_lines)
    
    with open(config_file, "w") as f:
        f.write(new_config)
    
    # 运行 MediaCrawler
    env = os.environ.copy()
    cookie = get_zhihu_cookie()
    if cookie:
        env["ZHIHU_COOKIE"] = cookie
    
    cmd = [
        "python3", "main.py",
        "--platform", "zhihu",
        "--lt", "cookie",
        "--type", "detail"
    ]
    
    print(f"运行命令: {' '.join(cmd)}")
    print(f"工作目录: {MEDIACRAWLER_DIR}")
    
    result = subprocess.run(
        cmd,
        cwd=MEDIACRAWLER_DIR,
        env=env,
        capture_output=True,
        text=True,
        timeout=120
    )
    
    print(f"返回码: {result.returncode}")
    if result.returncode != 0:
        print(f"stderr: {result.stderr[:500]}")
    
    # 查找最新的输出文件
    output_files = sorted(data_dir.glob("detail_contents_*.json"), reverse=True)
    
    if output_files:
        latest_file = output_files[0]
        with open(latest_file) as f:
            data = json.load(f)
        
        if data:
            item = data[0]
            return {
                "success": True,
                "title": item.get("title", ""),
                "content": item.get("content_text", ""),
                "url": item.get("content_url", ""),
                "author": item.get("user_nickname", ""),
                "created_time": item.get("created_time", 0),
                "voteup": item.get("voteup_count", 0),
                "comments": item.get("comment_count", 0),
                "content_type": item.get("content_type", "article"),
            }
    
    return {"success": False, "error": "未获取到数据"}


def save_to_markdown(data: dict, output_path: Path) -> str:
    """保存为 Markdown 格式"""
    created = datetime.fromtimestamp(data["created_time"]).strftime("%Y-%m-%d %H:%M:%S") if data.get("created_time") else ""
    
    content = f"""---
title: "{data['title']}"
author: "{data['author']}"
url: "{data['url']}"
created: "{created}"
voteup: {data.get('voteup', 0)}
comments: {data.get('comments', 0)}
platform: zhihu
type: article
---

# {data['title']}

**作者**: {data['author']}  
**来源**: [知乎专栏]({data['url']})  
**发布时间**: {created}  
**点赞**: {data.get('voteup', 0)}  
**评论**: {data.get('comments', 0)}

---

{data['content']}
"""
    
    with open(output_path, "w", encoding="utf-8") as f:
        f.write(content)
    
    return str(output_path)


def main():
    parser = argparse.ArgumentParser(description="知乎文章爬取工具 (基于 MediaCrawler)")
    parser.add_argument("--url", "-u", required=True, help="知乎文章链接")
    parser.add_argument("--output", "-o", default=str(DEFAULT_OUTPUT_DIR), help="输出目录")
    
    args = parser.parse_args()
    
    # 加载环境变量
    load_env()
    
    # 确保输出目录存在
    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    print(f"目标文章: {args.url}")
    print(f"输出目录: {output_dir}")
    
    # 爬取文章
    print("\n开始爬取...")
    result = crawl_zhihu_article(args.url, output_dir)
    
    if not result.get("success"):
        print(f"爬取失败: {result.get('error', '未知错误')}")
        sys.exit(1)
    
    print(f"\n爬取成功!")
    print(f"标题: {result['title']}")
    print(f"内容长度: {len(result['content'])} 字符")
    
    # 生成文件名
    safe_title = "".join(c for c in result['title'] if c.isalnum() or c in " -_")[:50]
    filename = f"zhihu_article_{safe_title}.md"
    output_path = output_dir / filename
    
    # 保存为 Markdown
    save_to_markdown(result, output_path)
    print(f"\n已保存到: {output_path}")
    
    # 输出 JSON 格式
    print("\n--- JSON 输出 ---")
    print(json.dumps({
        "success": True,
        "type": "article",
        "title": result["title"],
        "filepath": str(output_path),
        "content_length": len(result["content"]),
        "message": f"爬取成功: {result['title']}"
    }, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()