stock/local/backend/app/api/data_sync.py

"""数据同步 — 将本地 data/ 打包推送到 serve 端。"""
from __future__ import annotations

import json
import logging
import tarfile
import tempfile
from datetime import datetime
from io import BytesIO
from pathlib import Path

import httpx
from fastapi import APIRouter, HTTPException, Request
from pydantic import BaseModel

from app.config import settings

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/data/sync", tags=["data-sync"])

# 允许同步的子目录 — 覆盖 serve 功能面板所需的数据
SYNCABLE_PARTS = [
    "kline_daily",
    "kline_daily_enriched",
    "kline_index_daily",
    "kline_index_enriched",
    "kline_etf_daily",
    "kline_etf_enriched",
    "kline_etf_minute",
    "kline_minute",
    "adj_factor",
    "adj_factor_etf",
    "instruments",
    "instruments_index",
    "instruments_etf",
    "financials",
    "pools",
    "ext_data",
]


def _part_size_info(data_dir: Path, part: str) -> dict:
    """获取指定子目录的文件数和大小。"""
    part_dir = data_dir / part
    if not part_dir.exists():
        return {"file_count": 0, "size_bytes": 0, "last_modified": None}
    files = list(part_dir.rglob("*.parquet"))
    if not files:
        return {"file_count": 0, "size_bytes": 0, "last_modified": None}
    total = sum(f.stat().st_size for f in files)
    mtimes = [f.stat().st_mtime for f in files]
    last_ts = datetime.fromtimestamp(max(mtimes)).isoformat() if mtimes else None
    return {"file_count": len(files), "size_bytes": total, "last_modified": last_ts}


def _pack_parts(data_dir: Path, parts: list[str]) -> bytes:
    """将指定子目录打包为 tar.gz 字节流。"""
    buf = BytesIO()
    with tarfile.open(fileobj=buf, mode="w:gz") as tar:
        for part in parts:
            part_dir = data_dir / part
            if not part_dir.exists():
                continue
            for fpath in sorted(part_dir.rglob("*.parquet")):
                rel = fpath.relative_to(data_dir)
                tar.add(fpath, arcname=str(rel))
    buf.seek(0)
    return buf.getvalue()


class SyncConfig(BaseModel):
    serve_url: str = ""
    sync_key: str = ""
    enable_auto: bool = False
    interval_minutes: int = 0


# ── 存储配置用的 JSON 文件 ──────────────────────────────────
CONFIG_FILE = "sync_config.json"


def _config_path(data_dir: Path) -> Path:
    return data_dir / "user_data" / CONFIG_FILE


def _load_config(data_dir: Path) -> SyncConfig:
    cp = _config_path(data_dir)
    if cp.exists():
        try:
            return SyncConfig(**json.loads(cp.read_text()))
        except Exception:
            pass
    return SyncConfig()


def _save_config(data_dir: Path, cfg: SyncConfig) -> None:
    cp = _config_path(data_dir)
    cp.parent.mkdir(parents=True, exist_ok=True)
    cp.write_text(cfg.model_dump_json(indent=2))


# ── 路由 ──────────────────────────────────────────────────


@router.get("/status")
def sync_status(request: Request):
    """各数据目录的本地状态 + 同步配置。"""
    data_dir: Path = request.app.state.datastore.data_dir
    cfg = _load_config(data_dir)
    parts = {}
    for p in SYNCABLE_PARTS:
        parts[p] = _part_size_info(data_dir, p)
    return {
        "parts": parts,
        "config": cfg.model_dump(),
        "last_sync": cfg.model_dump().get("_last_sync_time"),
    }


class ConfigReq(BaseModel):
    serve_url: str
    sync_key: str
    enable_auto: bool = False
    interval_minutes: int = 0


@router.post("/config")
def update_sync_config(request: Request, body: ConfigReq):
    """更新同步配置。"""
    data_dir: Path = request.app.state.datastore.data_dir
    cfg = SyncConfig(
        serve_url=body.serve_url.rstrip("/"),
        sync_key=body.sync_key,
        enable_auto=body.enable_auto,
        interval_minutes=body.interval_minutes,
    )
    _save_config(data_dir, cfg)
    return {"ok": True}


class PushReq(BaseModel):
    parts: list[str]


@router.post("/push")
async def push_sync(request: Request, body: PushReq):
    """打包指定 parts 推送到 serve 端。"""
    data_dir: Path = request.app.state.datastore.data_dir
    cfg = _load_config(data_dir)
    if not cfg.serve_url:
        raise HTTPException(status_code=400, detail="未配置 serve_url，请先在同步设置中配置")

    requested = [p for p in body.parts if p in SYNCABLE_PARTS]
    if not requested:
        raise HTTPException(status_code=400, detail="未指定有效的同步目录")

    # 检查有没有数据可同步
    available = [p for p in requested if (data_dir / p).exists()]
    if not available:
        raise HTTPException(status_code=400, detail="所选目录均无数据可同步")

    # 打包
    tarball = _pack_parts(data_dir, available)
    size_bytes = len(tarball)
    logger.info("packed %d parts -> %d bytes", len(available), size_bytes)

    # 推送
    upload_url = f"{cfg.serve_url}/api/data/sync/upload"
    headers = {"X-Sync-Key": cfg.sync_key} if cfg.sync_key else {}
    try:
        async with httpx.AsyncClient(timeout=600) as client:
            resp = await client.post(
                upload_url,
                files={"file": ("sync.tar.gz", tarball, "application/gzip")},
                data={"parts": ",".join(available)},
                headers=headers,
            )
            if resp.status_code != 200:
                detail = resp.text[:500]
                raise HTTPException(
                    status_code=502,
                    detail=f"serve 返回异常 ({resp.status_code}): {detail}",
                )
            result = resp.json()
    except httpx.RequestError as exc:
        raise HTTPException(status_code=502, detail=f"连接 serve 失败: {exc}") from exc

    # 记录上次同步时间
    cfg_dict = cfg.model_dump()
    cfg_dict["_last_sync_time"] = datetime.now().isoformat()
    cfg = SyncConfig(**cfg_dict)
    _save_config(data_dir, cfg)

    logger.info("sync pushed: parts=%s files=%d", available, result.get("file_count", 0))
    return {
        "ok": True,
        "parts_sent": available,
        "size_bytes": size_bytes,
        "serve_response": result,
    }