InboundVerify/server.py

# server.py
#
# 服务模式入口：FastAPI（主线程，uvicorn/asyncio）+ Playwright worker（独立线程）。
# 客户端经 HTTP 触发任务、查状态、下载数据；worker 串行执行任务并跑心跳。
#
# 线程模型（关键）：
#   - 主线程：uvicorn + FastAPI。路由【绝不】访问 Playwright 对象，只经
#     task_queue（投递任务）+ state_store（查状态/任务）+ 文件系统（下载数据）。
#   - worker 线程：runtime.launch_and_prepare（sync_playwright 在此）+ 任务循环，
#     独占所有 page 操作；与主线程仅经 Queue + SQLite 通信。
#   违反"路由不碰 Playwright"会崩（sync 对象跨线程访问）。
#
# 运行：python server.py  （默认监听 0.0.0.0:8000）

import os
import queue
import threading
import time
from contextlib import asynccontextmanager
from typing import Dict, Optional

import uvicorn
from apscheduler.schedulers.background import BackgroundScheduler
from apscheduler.triggers.cron import CronTrigger
from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel

from paths import DOWNLOAD_DIR, OUTPUT_DIR
import state_store
from runtime import (
    HEARTBEAT_INTERVAL,
    TASK_HANDLERS,
    dispatch_task,
    launch_and_prepare,
    run_heartbeat,
)

# 全部站点；百世固定下载当天，不可配置偏移
ALL_SITES = ["顺心", "百世", "中通", "韵达", "安能"]
CONFIGURABLE_SITES = {"顺心", "中通", "韵达", "安能"}

# 任务队列：元素 (task_id, task_spec)。主线程投递，worker 消费。
task_queue: "queue.Queue" = queue.Queue()
# worker 运行状态（主线程只读，worker 写）
worker_state = {
    "ctx": None,
    "stop": False,
    "thread": None,
    "ready": False,  # launch_and_prepare 完成（各站就绪，可接任务）
    "error": None,  # worker 启动失败原因
}

# 每日定时下载调度器（进程内；job 只往 task_queue 投任务，不碰 Playwright）
scheduler = BackgroundScheduler(daemon=True)


def _worker_loop():
    """worker 线程：启动 Playwright + 等就绪 + 任务循环（执行任务 + 心跳）。"""
    try:
        ctx = launch_and_prepare()
        worker_state["ctx"] = ctx
        worker_state["ready"] = True
        print(">> [worker] 各站就绪，开始接收任务 ...")
    except Exception as e:
        worker_state["error"] = str(e)
        print(f"❌ [worker] 启动失败: {e}")
        return

    last_heartbeat = 0.0
    while not worker_state["stop"]:
        try:
            task_id, task_spec = task_queue.get(timeout=1)
        except queue.Empty:
            # 空闲时跑心跳
            if time.monotonic() - last_heartbeat >= HEARTBEAT_INTERVAL:
                run_heartbeat(ctx)
                last_heartbeat = time.monotonic()
            continue

        state_store.update_task(task_id, state_store.TASK_RUNNING)
        print(f">> [worker] 执行任务 #{task_id}: {task_spec}")
        status, error = dispatch_task(ctx, task_spec)
        state_store.update_task(task_id, status, error)
        print(f">> [worker] 任务 #{task_id} 完成: {status} {error or ''}")

    try:
        ctx.stop()
    except Exception:
        pass
    worker_state["ready"] = False
    print(">> [worker] 已退出。")


def _enqueue_undelivered(site):
    """定时 job：把该站 undelivered 任务投到队列（worker 串行处理；本线程不碰 Playwright）。"""
    try:
        tid = state_store.create_task(site, "undelivered")
        task_queue.put((tid, {"site": site, "kind": "undelivered"}))
        print(f">> [定时] 投递 {site}/undelivered 任务 #{tid}")
    except Exception as e:
        print(f">> [定时] 投递 {site} 失败: {e}")


def _reschedule_site(site):
    """按持久化配置（重新）注册或取消该站的每日定时 job。"""
    job_id = f"site_{site}"
    try:
        scheduler.remove_job(job_id)
    except Exception:
        pass
    cfg = state_store.get_all_config().get(site)
    if not cfg or not cfg.get("schedule_enabled") or not cfg.get("schedule_time"):
        return
    try:
        hh, mm = cfg["schedule_time"].split(":")
        scheduler.add_job(
            _enqueue_undelivered,
            CronTrigger(hour=int(hh), minute=int(mm)),
            args=[site],
            id=job_id,
            replace_existing=True,
        )
        print(f">> [定时] 已注册 {site} 每日 {cfg['schedule_time']} 下载")
    except Exception as e:
        print(f">> [定时] 注册 {site} 失败: {e}")


@asynccontextmanager
async def lifespan(_app):
    """服务启停：起 worker 线程 / 通知 worker 停。"""
    state_store.init_db()  # 先建表/迁移状态库，确保早于 worker 就绪的 /api/status 可用
    for site in ALL_SITES:  # 按持久化配置注册各站每日定时 job
        _reschedule_site(site)
    scheduler.start()
    print(">> [定时] 调度器已启动")
    t = threading.Thread(target=_worker_loop, daemon=True)
    worker_state["thread"] = t
    t.start()
    yield
    scheduler.shutdown(wait=False)
    print(">> [定时] 调度器已停止")
    worker_state["stop"] = True
    t.join(timeout=10)


app = FastAPI(title="InboundVerify 服务端", lifespan=lifespan)


class TaskRequest(BaseModel):
    site: str
    kind: str


@app.post("/tasks")
def create_task(req: TaskRequest):
    """提交任务 {site, kind} → 入队，返回 task_id。"""
    if (req.site, req.kind) not in TASK_HANDLERS:
        raise HTTPException(status_code=400, detail=f"无效任务: {req.site}/{req.kind}")
    task_id = state_store.create_task(req.site, req.kind)
    task_queue.put((task_id, {"site": req.site, "kind": req.kind}))
    return {"task_id": task_id}


@app.get("/tasks/{task_id}")
def get_task(task_id: int):
    t = state_store.get_task(task_id)
    if not t:
        raise HTTPException(status_code=404, detail="任务不存在")
    return t


@app.get("/tasks")
def list_tasks(limit: int = 20):
    return state_store.list_tasks(limit)


@app.get("/status")
def get_status():
    """各站登录态 + 数据态（前端状态盘用），另含 worker 就绪状态。"""
    return {
        "worker_ready": worker_state["ready"],
        "worker_error": worker_state["error"],
        "sites": state_store.get_all_status(),
    }


class ConfigRequest(BaseModel):
    expected_offset: Optional[int] = None
    actual_offset: Optional[int] = None
    schedule_enabled: Optional[bool] = None
    schedule_time: Optional[str] = None
    settings: Optional[Dict[str, str]] = None


def _default_cfg():
    return {
        "expected_offset": 0,
        "actual_offset": 0,
        "schedule_enabled": False,
        "schedule_time": "",
    }


@app.get("/config")
def get_config():
    """各站配置：应到/实到偏移 + 每日定时（百世偏移恒 0）。"""
    cfg = state_store.get_all_config()
    return {site: cfg.get(site, _default_cfg()) for site in ALL_SITES}


@app.put("/config/{site}")
def set_config(site: str, req: ConfigRequest):
    if site not in ALL_SITES:
        raise HTTPException(status_code=400, detail=f"未知站点: {site}")
    cur = state_store.get_all_config().get(site, _default_cfg())
    # 应到/实到偏移（百世锁定当天）
    for kind, val in (
        ("expected", req.expected_offset),
        ("actual", req.actual_offset),
    ):
        if val is not None:
            if site == "百世":
                raise HTTPException(
                    status_code=400, detail="百世固定下载当天，不可配置偏移"
                )
            state_store.set_offset(site, kind, val)
    # 每日定时（所有站可配，含百世）
    if req.schedule_enabled is not None or req.schedule_time is not None:
        enabled = (
            req.schedule_enabled
            if req.schedule_enabled is not None
            else cur["schedule_enabled"]
        )
        time_str = (
            req.schedule_time if req.schedule_time is not None else cur["schedule_time"]
        )
        state_store.set_schedule(site, enabled, time_str)
        _reschedule_site(site)
    # 站点专属配置（密码/账号/路径…）
    if req.settings:
        for k, v in req.settings.items():
            state_store.set_setting(site, k, v)
    cfg = state_store.get_all_config().get(site, _default_cfg())
    cfg["settings"] = state_store.get_site_settings(site)
    return cfg


@app.get("/config/{site}/settings")
def get_site_settings_api(site: str):
    """单站专属配置（密码/账号/路径…；不进 5s 轮询）。"""
    if site not in ALL_SITES:
        raise HTTPException(status_code=400, detail=f"未知站点: {site}")
    return state_store.get_site_settings(site)


REPORT_FILE = "应到未到数据.xlsx"


@app.get("/report")
def download_report():
    """下载 output/应到未到数据.xlsx（跑比对生成；未生成 404）。"""
    path = os.path.join(OUTPUT_DIR, REPORT_FILE)
    if not os.path.isfile(path):
        raise HTTPException(status_code=404, detail="报告尚未生成，请先跑比对")
    return FileResponse(path, filename=REPORT_FILE)


@app.get("/data/{filename}")
def download_data(filename: str):
    """下载 downloads/ 下的数据文件（防路径穿越）。"""
    if not filename or "/" in filename or "\\" in filename or ".." in filename:
        raise HTTPException(status_code=400, detail="非法文件名")
    path = os.path.join(DOWNLOAD_DIR, filename)
    # 双重校验：解析后绝对路径仍在 DOWNLOAD_DIR 内
    if not os.path.abspath(path).startswith(os.path.abspath(DOWNLOAD_DIR) + os.sep):
        raise HTTPException(status_code=400, detail="非法路径")
    if not os.path.isfile(path):
        raise HTTPException(status_code=404, detail="文件不存在")
    return FileResponse(path, filename=filename)


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)