xhs_crawler_claude/app/web.py

"""web.py - FastAPI 后端

目标：
- 提供基本的 Web UI + API（启动/停止采集、导入用户、登录/验证、日志、SSE）
- 单账号登录态长期复用（storage_state 保存）
- txt 导入：支持“数字小红书号”自动解析为 profile_id（并缓存，低频，减少风控）
- 移除无用/冗余的调试/自更新接口，代码保持清晰可维护
"""

from __future__ import annotations

import asyncio
import json
import os
from dataclasses import dataclass
from typing import List, Optional, Dict, Any

from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.responses import HTMLResponse, StreamingResponse, PlainTextResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi.requests import Request
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

from playwright.async_api import async_playwright, Browser

from app import config
from app.log_manager import LogManager
from app.scraper import XiaoHongShuScraper
from app import auth_manager
from app import user_resolver


# =========================
# App State
# =========================

@dataclass
class AppState:
    is_running: bool = False
    scraper: Optional[XiaoHongShuScraper] = None
    event_queue: Optional[asyncio.Queue] = None
    playwright: Any = None
    browser: Optional[Browser] = None


app_state = AppState(event_queue=asyncio.Queue())
log_manager = LogManager()


# =========================
# FastAPI init
# =========================

app = FastAPI(title="小红书数据采集器", version="1.1.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["*"],
)

# 静态资源 + 模板
app.mount("/static", StaticFiles(directory=os.path.join(config.BASE_DIR, "static")), name="static")
templates = Jinja2Templates(directory=os.path.join(config.BASE_DIR, "templates"))


@app.on_event("startup")
async def on_startup():
    """
    只做轻量初始化，不在启动时强制拉起浏览器。
    原因：
    - 很多环境首次运行未执行 `playwright install chromium`，如果启动即 launch，会导致 Web 直接起不来（你看到的“打不开”）。
    - 浏览器改为按需启动：第一次需要登录/解析/采集时再启动，失败也能在 UI 里看到明确错误。
    """
    app_state.playwright = None
    app_state.browser = None
    log_manager.append("INFO", "App started (browser will be launched lazily on demand)")


@app.on_event("shutdown")
async def on_shutdown():
    # 安全关闭
    try:
        if app_state.browser:
            await app_state.browser.close()
    finally:
        app_state.browser = None

    try:
        if app_state.playwright:
            await app_state.playwright.stop()
    finally:
        app_state.playwright = None


# =========================
# Models
# =========================

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


class StatusResponse(BaseModel):
    running: bool
    mode: Optional[str] = None
    user_count: int = 0


class UsersResponse(BaseModel):
    users: List[str]
    count: int


class ImportResponse(BaseModel):
    users: List[str]
    count: int
    total_lines: int
    imported: int
    skipped: int
    resolved: int
    failed_resolve: int
    duplicates: int
    invalid: int
    message: str
    # 可选返回一些失败明细（避免过大，只给前 20 条）
    failed_examples: List[str] = []


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


# =========================
# Pages
# =========================

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


# =========================
# Helpers
# =========================

def _users_file() -> str:
    return os.path.join(config.DATA_DIR, "users.txt")


def read_users_txt() -> List[str]:
    path = _users_file()
    if not os.path.exists(path):
        return []
    users: List[str] = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            users.append(line)
    # 去重但保持顺序
    seen = set()
    out = []
    for u in users:
        if u in seen:
            continue
        seen.add(u)
        out.append(u)
    return out


def write_users_txt(users: List[str]) -> None:
    path = _users_file()
    tmp = path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        for u in users:
            f.write(u + "\n")
    os.replace(tmp, path)


def require_browser() -> Browser:
    if app_state.browser is None:
        raise HTTPException(status_code=500, detail="Browser not ready")
    return app_state.browser


# =========================
# Status / Logs
# =========================

@app.get("/api/status", response_model=StatusResponse)
async def api_status():
    return StatusResponse(
        running=app_state.is_running,
        mode=getattr(app_state.scraper, "mode", None) if app_state.scraper else None,
        user_count=len(read_users_txt()),
    )


@app.post("/api/logs/clear")
async def clear_logs():
    log_manager.clear()
    return {"success": True}


@app.get("/api/logs")
async def get_logs(limit: int = 200):
    return {"logs": log_manager.get_latest(limit=limit)}


# =========================
# Users
# =========================

@app.get("/api/users", response_model=UsersResponse)
async def get_users():
    users = read_users_txt()
    return UsersResponse(users=users, count=len(users))


@app.post("/api/users")
async def add_user(user_id: str):
    user_id = (user_id or "").strip()
    if not user_id:
        raise HTTPException(status_code=400, detail="user_id 不能为空")

    users = read_users_txt()
    if user_id in users:
        return {"success": True, "message": "已存在"}

    users.append(user_id)
    write_users_txt(users)
    return {"success": True}


@app.delete("/api/users/{user_id}")
async def delete_user(user_id: str):
    users = read_users_txt()
    users = [u for u in users if u != user_id]
    write_users_txt(users)
    return {"success": True}


@app.delete("/api/users")
async def delete_all_users():
    write_users_txt([])
    return {"success": True}


@app.get("/api/users/export/txt")
async def export_users_txt():
    return PlainTextResponse("\\n".join(read_users_txt()))


@app.get("/api/users/import/txt", response_model=ImportResponse)
async def import_users_from_txt():
    """从 data/users.txt 读取并规范化：
    - 若是 profile_id / profile URL：直接收录
    - 若是纯数字小红书号：通过搜索页解析成 profile_id（Playwright，低频、缓存）
    - 输出：users.txt 最终仅保存 profile_id（稳定可采集）
    """
    path = _users_file()
    if not os.path.exists(path):
        raise HTTPException(status_code=400, detail="data/users.txt 不存在")

    with open(path, "r", encoding="utf-8") as f:
        raw_lines = [l.rstrip("\n") for l in f]

    total_lines = len(raw_lines)

    browser = require_browser()

    imported_ids: List[str] = []
    duplicates = 0
    invalid = 0
    skipped = 0
    resolved = 0
    failed_resolve = 0
    failed_examples: List[str] = []

    seen = set()

    # 第一轮：先把能直接解析的收集，数字 keyword 单独列表
    keywords: List[tuple[str, str]] = []  # (keyword, raw_line)
    for line in raw_lines:
        kw, pid, _label = user_resolver.parse_line_to_keyword_or_profile(line)
        if pid:
            pid = pid.lower()
            if pid in seen:
                duplicates += 1
                continue
            seen.add(pid)
            imported_ids.append(pid)
            continue
        if kw:
            keywords.append((kw, line))
        else:
            # 空行/注释等视为 skipped；其它无法识别为 invalid
            if not line.strip() or line.strip().startswith("#") or line.strip().startswith("//"):
                skipped += 1
            else:
                invalid += 1
                if len(failed_examples) < 20:
                    failed_examples.append(f"invalid: {line}")
            continue

    # 第二轮：逐个解析 keyword（单线程+慢一点，减少风控）
    for kw, raw in keywords:
        pid, msg = await user_resolver.resolve_keyword_to_profile_id(browser, kw, slow_mode=True)
        if not pid:
            failed_resolve += 1
            if len(failed_examples) < 20:
                failed_examples.append(f"resolve_failed({kw}): {raw} ({msg})")
            continue
        pid = pid.lower()
        if pid in seen:
            duplicates += 1
            continue
        seen.add(pid)
        imported_ids.append(pid)
        resolved += 1

        # 每解析一个稍微间隔一下（更像真人）
        await asyncio.sleep(0.8)

    # 写回：只保存 profile_id
    write_users_txt(imported_ids)

    imported = len(imported_ids)
    message = f"导入完成：{imported} 个 profile_id（解析成功 {resolved}）"

    log_manager.append("INFO", message)

    return ImportResponse(
        users=imported_ids,
        count=imported,
        total_lines=total_lines,
        imported=imported,
        skipped=skipped,
        resolved=resolved,
        failed_resolve=failed_resolve,
        duplicates=duplicates,
        invalid=invalid,
        message=message,
        failed_examples=failed_examples,
    )


# =========================
# Login
# =========================

@app.get("/api/login/status", response_model=LoginStatusResponse)
async def get_login_status():
    status = await auth_manager.get_login_status(app_state.browser)
    return LoginStatusResponse(
        logged_in=status.logged_in,
        login_time=status.login_time,
        last_verified_time=status.last_verified_time,
        message=status.message,
    )


@app.post("/api/login/qrcode")
async def open_login():
    """兼容旧前端：返回 login_url，同时在服务端打开 Playwright 登录页。

    说明：前端 window.open 打开的页面不会影响服务端 Playwright 的登录态。
    正确使用方式：在部署环境（Docker + VNC）里打开 VNC 浏览器扫码登录。
    """
    browser = require_browser()
    await auth_manager.open_login_page(browser)
    return {
        "login_url": config.LOGIN_URL,
        "hint": "已在服务端打开登录页。请在 VNC/有界面环境里扫码登录，然后点击“验证登录”。"
    }


@app.post("/api/login/verify")
async def verify_login():
    browser = require_browser()
    ok, msg = await auth_manager.verify_login(browser)
    meta = auth_manager.load_auth_meta()
    return {
        "success": ok,
        "message": None if ok else msg,
        "login_time": meta.get("login_time"),
        "expire_time": None,  # 单账号长期复用：不强制本地过期时间
        "last_verified_time": meta.get("last_verified_time"),
    }


@app.post("/api/logout")
async def logout():
    auth_manager.clear_login_state()
    return {"success": True}


# =========================
# Scrape control
# =========================

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

    users = read_users_txt()
    if not users:
        raise HTTPException(status_code=400, detail="用户列表为空，请先导入 users.txt")

    # 启动前校验登录态（单账号）
    browser = require_browser()
    ok, msg = await auth_manager.verify_login(browser)
    if not ok:
        raise HTTPException(status_code=400, detail=f"未登录或登录失效：{msg}")

    app_state.is_running = True
    app_state.scraper = XiaoHongShuScraper(app_state.event_queue, log_manager, browser)

    async def run_scraper():
        try:
            await app_state.scraper.run(user_ids=users, mode=request.mode, max_threads=request.max_threads)
        finally:
            app_state.is_running = False

    background_tasks.add_task(run_scraper)
    return {"success": True}


@app.post("/api/stop")
async def stop_scraping():
    if not app_state.is_running or not app_state.scraper:
        return {"success": True, "message": "未在运行"}

    await app_state.scraper.stop()
    app_state.is_running = False
    return {"success": True}


# =========================
# Data
# =========================

@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="数据不存在")
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# =========================
# SSE events
# =========================

@app.get("/api/events")
async def sse_events():
    async def event_generator():
        while True:
            event = await app_state.event_queue.get()
            yield f"data: {json.dumps(event, ensure_ascii=False)}\n\n"

    return StreamingResponse(event_generator(), media_type="text/event-stream")