xhs_crawler_claude/app/web.py

"""
Web Backend - FastAPI 应用
提供 API 和 SSE 事件推送
"""
import asyncio
import json
import os
import subprocess
from datetime import datetime
from typing import Dict, List, Optional
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi.requests import Request
from pydantic import BaseModel
from typing import List as ListType

from app import config
from app.log_manager import get_log_manager
from app.scraper import XiaoHongShuScraper


# ========== Pydantic 模型 ==========
class StartRequest(BaseModel):
    mode: str = "single"  # single / multi
    max_threads: int = 3


class StatusResponse(BaseModel):
    is_running: bool
    mode: str
    max_threads: int
    total_users: int
    completed_users: int
    failed_users: int
    active_threads: int
    users_status: Dict[str, Dict]


# ========== 全局状态 ==========
class AppState:
    def __init__(self):
        self.is_running = False
        self.mode = "single"
        self.max_threads = 1
        self.total_users = 0
        self.completed_users = 0
        self.failed_users = 0
        self.active_threads = 0
        self.users_status: Dict[str, Dict] = {}
        self.thread_status: Dict[str, Dict] = {}
        self.event_queue: asyncio.Queue = None
        self.scraper: Optional[XiaoHongShuScraper] = None
        self.scraper_task: Optional[asyncio.Task] = None


app_state = AppState()
log_manager = get_log_manager()

# ========== FastAPI 应用 ==========
app = FastAPI(title="小红书数据采集器", version="1.0.0")

# 挂载静态文件和模板
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")


# ========== 路由 ==========

@app.get("/", response_class=HTMLResponse)
async def index(request: Request):
    """首页"""
    return templates.TemplateResponse("index.html", {"request": request})




@app.get("/api/status", response_model=StatusResponse)
async def get_status():
    """获取当前状态"""
    return StatusResponse(
        is_running=app_state.is_running,
        mode=app_state.mode,
        max_threads=app_state.max_threads,
        total_users=app_state.total_users,
        completed_users=app_state.completed_users,
        failed_users=app_state.failed_users,
        active_threads=app_state.active_threads,
        users_status=app_state.users_status
    )


@app.post("/api/start")
async def start_scraping(request: StartRequest):
    """启动采集"""
    if app_state.is_running:
        raise HTTPException(status_code=400, detail="采集任务正在运行中")

    # 读取用户 ID 列表
    users_file = os.path.join(config.DATA_DIR, "users.txt")
    if not os.path.exists(users_file):
        raise HTTPException(status_code=400, detail="用户文件不存在，请先创建 data/users.txt")

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

    if not user_ids:
        raise HTTPException(status_code=400, detail="用户 ID 列表为空")

    # 验证模式
    if request.mode not in ["single", "multi"]:
        raise HTTPException(status_code=400, detail="无效的模式，必须是 single 或 multi")

    # 验证线程数
    if request.mode == "multi":
        if request.max_threads < 1 or request.max_threads > config.MAX_CONCURRENT_THREADS:
            raise HTTPException(status_code=400, detail=f"线程数必须在 1-{config.MAX_CONCURRENT_THREADS} 之间")

    # 重置状态
    app_state.is_running = True
    app_state.mode = request.mode
    app_state.max_threads = request.max_threads if request.mode == "multi" else 1
    app_state.total_users = len(user_ids)
    app_state.completed_users = 0
    app_state.failed_users = 0
    app_state.active_threads = 0
    app_state.users_status = {}
    app_state.thread_status = {}

    # 初始化用户状态
    for user_id in user_ids:
        app_state.users_status[user_id] = {
            "status": "pending",
            "note_count": 0,
            "thread_id": None
        }

    # 创建事件队列
    app_state.event_queue = asyncio.Queue()

    # 创建采集器
    app_state.scraper = XiaoHongShuScraper(app_state.event_queue, log_manager)

    # 启动采集任务
    async def run_scraper():
        try:
            await app_state.scraper.run(
                user_ids=user_ids,
                mode=request.mode,
                max_threads=request.max_threads
            )
        finally:
            app_state.is_running = False
            await app_state.event_queue.put({
                "type": "finished"
            })

    app_state.scraper_task = asyncio.create_task(run_scraper())

    return {
        "success": True,
        "message": f"采集任务已启动，模式: {request.mode}, 用户数: {len(user_ids)}",
        "total_users": len(user_ids),
        "mode": request.mode,
        "max_threads": app_state.max_threads
    }


@app.post("/api/stop")
async def stop_scraping():
    """停止采集"""
    if not app_state.is_running:
        raise HTTPException(status_code=400, detail="没有正在运行的采集任务")

    if app_state.scraper:
        app_state.scraper.stop()

    # 取消任务
    if app_state.scraper_task:
        app_state.scraper_task.cancel()
        try:
            await app_state.scraper_task
        except asyncio.CancelledError:
            pass

    app_state.is_running = False

    return {
        "success": True,
        "message": "采集任务已停止"
    }


@app.post("/api/logs/clear")
async def clear_logs():
    """清空日志"""
    log_manager.clear()
    return {
        "success": True,
        "message": "日志已清空"
    }


@app.get("/api/logs")
async def get_logs(level: Optional[str] = None, limit: Optional[int] = None):
    """获取日志"""
    logs = log_manager.get_logs(level=level, limit=limit)
    return {
        "logs": logs,
        "stats": log_manager.get_stats()
    }


@app.get("/api/logs/search")
async def search_logs(keyword: str, limit: Optional[int] = None):
    """搜索日志"""
    logs = log_manager.search_logs(keyword, limit)
    return {
        "logs": logs,
        "count": len(logs)
    }


@app.get("/api/logs/download")
async def download_logs():
    """下载日志文件"""
    text = log_manager.export_to_text()

    return StreamingResponse(
        iter([text]),
        media_type="text/plain",
        headers={
            "Content-Disposition": f"attachment; filename=xiaohongshu_logs_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
        }
    )


@app.get("/api/users")
async def get_users():
    """获取用户列表"""
    users_file = os.path.join(config.DATA_DIR, "users.txt")

    if not os.path.exists(users_file):
        return {"users": []}

    with open(users_file, 'r', encoding='utf-8') as f:
        users = [{"id": line.strip()} for line in f if line.strip()]

    # 添加采集状态
    for user in users:
        user_id = user["id"]
        if user_id in app_state.users_status:
            user.update(app_state.users_status[user_id])
        else:
            user["status"] = "pending"
            user["note_count"] = 0

    return {"users": users}


@app.post("/api/users")
async def add_user(request: dict):
    """添加用户"""
    user_id = request.get("user_id", "").strip()

    if not user_id:
        raise HTTPException(status_code=400, detail="用户ID不能为空")

    # 验证用户ID格式
    if not user_id or len(user_id) < 10:
        raise HTTPException(status_code=400, detail="用户ID格式不正确")

    users_file = os.path.join(config.DATA_DIR, "users.txt")

    # 检查是否已存在
    if os.path.exists(users_file):
        with open(users_file, 'r', encoding='utf-8') as f:
            existing_users = [line.strip() for line in f if line.strip()]
            if user_id in existing_users:
                raise HTTPException(status_code=400, detail="用户ID已存在")

    # 添加到文件
    with open(users_file, 'a', encoding='utf-8') as f:
        f.write(f"{user_id}\n")

    return {
        "success": True,
        "message": f"用户 {user_id} 已添加",
        "user_id": user_id
    }


@app.delete("/api/users/{user_id}")
async def remove_user(user_id: str):
    """删除用户"""
    users_file = os.path.join(config.DATA_DIR, "users.txt")

    if not os.path.exists(users_file):
        raise HTTPException(status_code=404, detail="用户文件不存在")

    # 读取所有用户
    with open(users_file, 'r', encoding='utf-8') as f:
        users = [line.strip() for line in f if line.strip()]

    # 检查用户是否存在
    if user_id not in users:
        raise HTTPException(status_code=404, detail="用户不存在")

    # 删除用户
    users.remove(user_id)

    # 写回文件
    with open(users_file, 'w', encoding='utf-8') as f:
        for user in users:
            f.write(f"{user}\n")

    # 删除采集数据文件
    data_file = os.path.join(config.DATA_DIR, f"{user_id}.json")
    if os.path.exists(data_file):
        os.remove(data_file)

    return {
        "success": True,
        "message": f"用户 {user_id} 已删除"
    }


@app.delete("/api/users")
async def clear_all_users():
    """清空所有用户"""
    users_file = os.path.join(config.DATA_DIR, "users.txt")

    if not os.path.exists(users_file):
        raise HTTPException(status_code=404, detail="用户文件不存在")

    # 清空文件
    with open(users_file, 'w', encoding='utf-8') as f:
        f.write("")

    # 删除所有采集数据
    data_dir = config.DATA_DIR
    for filename in os.listdir(data_dir):
        if filename.endswith('.json') and filename != 'auth_state.json':
            file_path = os.path.join(data_dir, filename)
            os.remove(file_path)

    # 重置状态
    app_state.users_status = {}

    return {
        "success": True,
        "message": "所有用户已清空"
    }


@app.post("/api/users/import/txt")
async def import_users_from_txt():
    """从users.txt导入用户（模拟）"""
    users_file = os.path.join(config.DATA_DIR, "users.txt")

    if not os.path.exists(users_file):
        raise HTTPException(status_code=404, detail="users.txt文件不存在")

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

    return {
        "success": True,
        "count": len(user_ids),
        "users": user_ids
    }


@app.get("/api/users/export/txt")
async def export_users_to_txt():
    """导出用户列表到txt文件"""
    users_file = os.path.join(config.DATA_DIR, "users.txt")

    if not os.path.exists(users_file):
        raise HTTPException(status_code=404, detail="用户文件不存在")

    with open(users_file, 'r', encoding='utf-8') as f:
        content = f.read()

    return StreamingResponse(
        iter([content]),
        media_type="text/plain",
        headers={
            "Content-Disposition": f"attachment; filename=users_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
        }
    )


# ========== 登录相关 API ==========

class LoginStatusResponse(BaseModel):
    logged_in: bool
    login_time: Optional[str] = None
    expire_time: Optional[str] = None


@app.get("/api/login/status", response_model=LoginStatusResponse)
async def get_login_status():
    """获取登录状态"""
    auth_file = os.path.join(config.DATA_DIR, "auth_state.json")

    if not os.path.exists(auth_file):
        return LoginStatusResponse(logged_in=False)

    try:
        with open(auth_file, 'r', encoding='utf-8') as f:
            auth_data = json.load(f)

        # 检查是否过期
        login_time = auth_data.get("login_time")
        expire_time = auth_data.get("expire_time")
        current_time = datetime.now().isoformat()

        if expire_time and current_time > expire_time:
            # 登录已过期，清除状态
            os.remove(auth_file)
            return LoginStatusResponse(logged_in=False)

        return LoginStatusResponse(
            logged_in=auth_data.get("logged_in", False),
            login_time=login_time,
            expire_time=expire_time
        )
    except Exception:
        return LoginStatusResponse(logged_in=False)


@app.post("/api/login/qrcode")
async def generate_qrcode():
    """模拟登录成功"""
    import uuid
    from datetime import datetime, timedelta

    # 生成登录状态（模拟登录成功）
    login_id = str(uuid.uuid4())
    current_time = datetime.now()
    expire_time = current_time + timedelta(hours=24)  # 24小时后过期

    # 保存登录状态到文件
    auth_file = os.path.join(config.DATA_DIR, "auth_state.json")
    auth_data = {
        "logged_in": True,
        "login_time": current_time.isoformat(),
        "expire_time": expire_time.isoformat(),
        "login_id": login_id,
        "platform": "xiaohongshu"
    }

    with open(auth_file, 'w', encoding='utf-8') as f:
        json.dump(auth_data, f, ensure_ascii=False, indent=2)

    # 返回登录信息
    return {
        "login_id": login_id,
        "login_time": current_time.isoformat(),
        "expire_time": expire_time.isoformat(),
        "message": "登录成功（模拟）",
        "expire_in_hours": 24
    }


@app.post("/api/logout")
async def logout():
    """退出登录"""
    auth_file = os.path.join(config.DATA_DIR, "auth_state.json")

    try:
        if os.path.exists(auth_file):
            os.remove(auth_file)

        return {
            "success": True,
            "message": "已退出登录"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"退出登录失败: {str(e)}")


@app.get("/api/data/{user_id}")
async def get_user_data(user_id: str):
    """获取用户采集数据"""
    file_path = os.path.join(config.DATA_DIR, f"{user_id}.json")

    if not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail="用户数据不存在")

    with open(file_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    return data


@app.get("/api/events")
async def event_stream():
    """SSE 事件流"""
    async def event_generator():
        while True:
            try:
                # 等待事件，超时 1 秒发送心跳
                event = await asyncio.wait_for(
                    app_state.event_queue.get(),
                    timeout=1.0
                )

                # 更新状态
                if event["type"] == "log":
                    # 日志不需要更新状态
                    pass
                elif event["type"] == "progress":
                    app_state.completed_users = event.get("completed", app_state.completed_users)
                    app_state.failed_users = event.get("failed", app_state.failed_users)
                    app_state.active_threads = event.get("active_threads", app_state.active_threads)
                elif event["type"] == "user_status":
                    if event["user_id"] in app_state.users_status:
                        app_state.users_status[event["user_id"]]["status"] = event["status"]
                        app_state.users_status[event["user_id"]]["note_count"] = event.get("note_count", 0)
                        app_state.users_status[event["user_id"]]["thread_id"] = event.get("thread_id")
                elif event["type"] == "thread_status":
                    app_state.thread_status[event["thread_id"]] = {
                        "status": event["status"],
                        "current_user": event.get("current_user")
                    }
                    # 更新活跃线程数
                    app_state.active_threads = sum(
                        1 for s in app_state.thread_status.values() if s.get("status") == "busy"
                    )

                # 发送事件
                yield f"data: {json.dumps(event)}\n\n"

            except asyncio.TimeoutError:
                # 发送心跳
                yield ": heartbeat\n\n"

            except Exception as e:
                # 错误事件
                yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Access-Control-Allow-Origin": "*"
        }
    )


# ========== Git 相关 API ==========

@app.get("/api/git/status")
async def git_status():
    """获取 Git 状态"""
    if not os.path.exists("/app/.git"):
        return {
            "enabled": False,
            "message": "不是 Git 仓库"
        }

    try:
        # 获取当前分支
        result = subprocess.run(
            ["git", "branch", "--show-current"],
            capture_output=True,
            text=True,
            cwd="/app"
        )
        current_branch = result.stdout.strip()

        # 获取最新提交
        result = subprocess.run(
            ["git", "log", "-1", "--format=%h %s %cr"],
            capture_output=True,
            text=True,
            cwd="/app"
        )
        latest_commit = result.stdout.strip()

        # 检查是否有未提交的更改
        result = subprocess.run(
            ["git", "status", "--porcelain"],
            capture_output=True,
            text=True,
            cwd="/app"
        )
        has_changes = len(result.stdout.strip()) > 0

        return {
            "enabled": True,
            "branch": current_branch,
            "latest_commit": latest_commit,
            "has_changes": has_changes,
            "auto_pull_enabled": config.ENABLE_GIT_PULL,
            "auto_pull_interval": config.AUTO_PULL_INTERVAL
        }

    except Exception as e:
        return {
            "enabled": False,
            "error": str(e)
        }


@app.post("/api/git/pull")
async def git_pull():
    """手动拉取最新代码"""
    if not os.path.exists("/app/.git"):
        raise HTTPException(status_code=400, detail="不是 Git 仓库")

    try:
        # 获取当前分支
        result = subprocess.run(
            ["git", "branch", "--show-current"],
            capture_output=True,
            text=True,
            cwd="/app"
        )
        current_branch = result.stdout.strip()

        # 拉取最新代码
        result = subprocess.run(
            ["git", "fetch", "origin", current_branch],
            capture_output=True,
            text=True,
            cwd="/app"
        )

        if result.returncode != 0:
            raise HTTPException(status_code=500, detail=f"拉取失败: {result.stderr}")

        # 检查是否有更新
        result = subprocess.run(
            ["git", "rev-list", "HEAD..origin/" + current_branch, "--count"],
            capture_output=True,
            text=True,
            cwd="/app"
        )
        commits_behind = int(result.stdout.strip())

        if commits_behind == 0:
            return {
                "success": True,
                "message": "已是最新代码",
                "needs_restart": False
            }

        # 重置到最新代码
        result = subprocess.run(
            ["git", "reset", "--hard", "origin/" + current_branch],
            capture_output=True,
            text=True,
            cwd="/app"
        )

        if result.returncode != 0:
            raise HTTPException(status_code=500, detail=f"重置失败: {result.stderr}")

        # 检查 Python 文件是否有变化
        result = subprocess.run(
            ["git", "diff", "--name-only", "HEAD@{1}", "HEAD"],
            capture_output=True,
            text=True,
            cwd="/app"
        )
        changed_files = result.stdout.strip().split('\n')
        python_changed = any(f.endswith('.py') for f in changed_files if f)

        return {
            "success": True,
            "message": f"已拉取 {commits_behind} 个新提交",
            "changed_files": changed_files,
            "needs_restart": python_changed,
            "restart_hint": "Python 代码有变化，建议重启容器以应用更新"
        }

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"拉取失败: {str(e)}")


@app.post("/api/git/restart")
async def restart_container():
    """
    请求重启容器
    注意：这需要在容器外配置重启策略
    """
    return {
        "success": True,
        "message": "重启请求已接收",
        "methods": [
            "docker restart xiaohongshu-scraper",
            "docker-compose restart"
        ]
    }


# ========== 启动事件 ==========
@app.on_event("startup")
async def startup_event():
    """启动时初始化"""
    # 确保数据目录存在
    os.makedirs(config.DATA_DIR, exist_ok=True)
    os.makedirs(config.SCREENSHOT_DIR, exist_ok=True)

    # 初始化事件队列
    app_state.event_queue = asyncio.Queue()

    print("=" * 60)
    print("小红书数据采集器 Web 服务启动")
    print(f"访问地址: http://{config.WEB_HOST}:{config.WEB_PORT}")
    print("=" * 60)


@app.on_event("shutdown")
async def shutdown_event():
    """关闭时清理"""
    if app_state.scraper:
        app_state.scraper.stop()

    if app_state.scraper_task:
        app_state.scraper_task.cancel()

    print("Web 服务已关闭")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        "app.web:app",
        host=config.WEB_HOST,
        port=config.WEB_PORT,
        reload=True
    )