workbuddy-portal/manage.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""WorkBuddy Portal —— 统一命令行入口。

采集 / 存储 / 呈现三件事都由本项目承担，不再依赖外部计划任务或自动化。

常用：
  python manage.py init                       初始化数据库（建表 + 默认配置 + 管理员）
  python manage.py serve                      启动 Web（0.0.0.0:8848，进程内含调度线程）
  python manage.py serve --port 9000 --debug  开发模式（reloader 下调度只启动一份）
  python manage.py collect                    执行一次增量采集并退出（可用于外部计划任务）
  python manage.py migrate-csv [文件]          从旧版 CSV 存档导入（默认自动探测路径）
  python manage.py import-xlsx <文件>          从官网导出的 xlsx 合入
  python manage.py fill-prompt                补全缺失的 User Prompt
  python manage.py export-csv [路径]           导出与官网同构的 CSV
  python manage.py stats                      只看存档概况，不联网
  python manage.py passwd <用户名> [新密码]     重置登录密码
  python manage.py status                     查看调度与最近采集状态
"""
import argparse
import json
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from workbuddy_portal import client, collect, config, db, query, scheduler   # noqa: E402


def _p(*a):
    print(*a)


def cmd_init(args):
    db.init_db(admin_user=args.user, admin_password=args.password)
    conn = db.connect()
    try:
        n = query.totals(conn)
        _p("数据库已就绪：%s" % config.SQLITE_PATH)
        _p("  存档 %d 条 / %.2f 积分 / %d 个活跃日" % (n["records"], n["credits"], n["days"]))
        _p("  管理员：%s" % args.user)
    finally:
        conn.close()


def cmd_serve(args):
    from workbuddy_portal import create_app
    app = create_app(start_scheduler=not args.no_scheduler)
    host = args.host or config.DEFAULT_HOST
    port = args.port or config.DEFAULT_PORT
    if args.debug:
        app.run(host=host, port=port, debug=True)
        return
    try:
        from waitress import serve
        _p("生产模式（waitress）监听 http://%s:%d" % (host, port))
        serve(app, host=host, port=port, threads=8, ident="workbuddy-portal")
    except ImportError:
        _p("[warn] 未安装 waitress，回退到 Flask 内置服务器（生产建议 pip install waitress）")
        app.run(host=host, port=port, threaded=True)


def cmd_collect(args):
    db.init_db(create_admin=False)
    try:
        r = collect.run_sync(trigger="cli")
    except collect.Busy as e:
        _p("[busy] %s" % e)
        return 1
    except collect.ApiError as e:
        _p("[error] %s" % e)
        return 3 if e.cookie_expired else 5
    for line in r["lines"]:
        _p(line)
    return 0


def cmd_migrate_csv(args):
    db.init_db(create_admin=False)
    path = args.path
    if not path:
        for c in config.LEGACY_CSV_CANDIDATES:
            if os.path.exists(c):
                path = c
                break
    if not path:
        _p("[error] 找不到旧存档 CSV，请显式指定路径")
        return 2
    conn = db.connect()
    try:
        collect.migrate_from_csv(conn, path, log=_p)
        n = query.totals(conn)
        _p("当前存档：%d 条 / %.2f 积分 / %s ~ %s" % (n["records"], n["credits"],
                                                    n["firstDay"], n["lastDay"]))
    finally:
        conn.close()
    return 0


def cmd_import_xlsx(args):
    db.init_db(create_admin=False)
    conn = db.connect()
    try:
        collect.import_xlsx(conn, args.path, log=_p)
    finally:
        conn.close()


def cmd_fill_prompt(args):
    db.init_db(create_admin=False)
    conn = db.connect()
    try:
        collect.fill_prompt(conn, log=_p)
    finally:
        conn.close()


def cmd_export_csv(args):
    db.init_db(create_admin=False)
    conn = db.connect()
    try:
        path, n = collect.export_csv(conn, args.path)
        _p("已导出 %d 条 -> %s" % (n, path))
    finally:
        conn.close()


def cmd_stats(args):
    db.init_db(create_admin=False)
    conn = db.connect()
    try:
        t = query.totals(conn)
        if not t["records"]:
            _p("存档为空，先跑 python manage.py migrate-csv 或 manage.py collect")
            return
        _p("存档：%d 条 / %.2f 积分 / %d 个活跃日（%s ~ %s）"
           % (t["records"], t["credits"], t["days"], t["firstDay"], t["lastDay"]))
        _p("计费调用 %d · 免费调用 %d · 模型 %d · 客户端 %d"
           % (t["billableCalls"], t["freeCalls"], t["models"], t["clients"]))
        _p("")
        _p("%-24s %8s %12s %10s %8s" % ("模型", "调用", "积分", "单次均价", "免费占比"))
        for m in query.dims(conn)["model"]:
            _p("%-24s %8d %12.2f %10.2f %7.0f%%"
               % (m["name"], m["calls"], m["credits"], m["avgPerCall"], m["freeRate"] * 100))
        runs = conn.execute("SELECT id,trigger,status,started_at,added,dup,total,message"
                            " FROM collect_runs ORDER BY id DESC LIMIT 5").fetchall()
        if runs:
            _p("")
            _p("最近采集：")
            for r in runs:
                _p("  #%d %s %s +%d/%d → %d  %s"
                   % (r["id"], r["started_at"], r["status"], r["added"], r["dup"],
                      r["total"], r["message"] or ""))
    finally:
        conn.close()


def cmd_import_creds(args):
    """把 VSCode 设置里的 cookie / userAgent 接管进数据库（一次性迁移用）。"""
    db.init_db(create_admin=False)
    conn = db.connect()
    try:
        found = client.read_vscode_creds()
        if not found:
            _p("[error] 没找到 VSCode 系编辑器的 settings.json")
            return 2
        hit = None
        for path, cookie, ua, note in found:
            _p("  %-70s %s" % (path, note))
            if cookie and hit is None:
                hit = (cookie, ua)
        if not hit:
            _p("[error] 这些文件里都没有 codebuddyUsage.cookie，请到「配置管理」页手工粘贴")
            return 2
        cookie, ua = hit
        db.set_setting(conn, "cookie", cookie)
        if ua:
            db.set_setting(conn, "user_agent", ua)
        db.audit(conn, "import_creds", "cli", "从 VSCode 设置导入凭证（%d 字符）" % len(cookie), "127.0.0.1")
        _p("已导入 Cookie（%d 字符）与 User-Agent（%s）" % (len(cookie), "有" if ua else "无"))
    finally:
        conn.close()
    return 0


def cmd_passwd(args):
    db.init_db(create_admin=False)
    from workbuddy_portal.security import hash_password
    conn = db.connect()
    try:
        row = conn.execute("SELECT id FROM users WHERE username=?", (args.user,)).fetchone()
        pwd = args.password or "admin123"
        if row:
            conn.execute("UPDATE users SET password_hash=? WHERE id=?", (hash_password(pwd), row["id"]))
            _p("已重置 %s 的密码" % args.user)
        else:
            conn.execute("INSERT INTO users(username,password_hash,display_name,is_admin,created_at)"
                         " VALUES(?,?,?,1,?)", (args.user, hash_password(pwd), args.user, db.now_str()))
            _p("已创建用户 %s" % args.user)
        _p("新密码：%s" % pwd)
    finally:
        conn.close()


def cmd_status(args):
    db.init_db(create_admin=False)
    conn = db.connect()
    try:
        _p("服务器时间：%s" % db.now_str())
        _p("调度开关：%s" % ("启用" if db.get_bool(conn, "schedule_enabled", True) else "停用"))
        _p("每日时刻：%s" % (", ".join(scheduler.slots(conn)) or "—"))
        nxt = scheduler.next_run_at(conn)
        _p("下次执行：%s" % (nxt.strftime("%Y-%m-%d %H:%M:%S") if nxt else "—"))
        _p("Cookie：%s" % ("已配置" if (db.get_setting(conn, "cookie") or "").strip() else "未配置"))
        _p("互斥锁：%s" % ("存在（有采集在跑）" if os.path.exists(collect.LOCK_PATH) else "不存在"))
        last = conn.execute("SELECT * FROM collect_runs ORDER BY id DESC LIMIT 1").fetchone()
        if last:
            _p("最近采集：#%d %s %s %s" % (last["id"], last["started_at"], last["status"],
                                        last["message"] or ""))
        else:
            _p("最近采集：无")
        _p("（注意：调度线程只在 manage.py serve 进程内运行）")
    finally:
        conn.close()


def cmd_vacuum(args):
    """整理数据库：checkpoint WAL + VACUUM 回收空间。"""
    db.init_db(create_admin=False)
    conn = db.connect()
    try:
        before = os.path.getsize(config.SQLITE_PATH) if os.path.exists(config.SQLITE_PATH) else 0
        conn.execute("PRAGMA wal_checkpoint(TRUNCATE)")
        conn.execute("VACUUM")
        conn.execute("PRAGMA optimize")
        after = os.path.getsize(config.SQLITE_PATH) if os.path.exists(config.SQLITE_PATH) else 0
        _p("数据库整理完成：%s → %s（%+d 字节）" % (_human(before), _human(after), after - before))
        _p("存档 %d 条记录" % collect.record_count(conn))
    finally:
        conn.close()
    return 0


def _human(n):
    for unit in ("B", "KB", "MB", "GB"):
        if n < 1024 or unit == "GB":
            return ("%d B" % n) if unit == "B" else ("%.1f %s" % (n, unit))
        n /= 1024.0


def main():
    ap = argparse.ArgumentParser(description="WorkBuddy Portal（workbuddy-portal）",
                                 formatter_class=argparse.RawDescriptionHelpFormatter,
                                 epilog=__doc__)
    sub = ap.add_subparsers(dest="cmd")

    s = sub.add_parser("init", help="初始化数据库")
    s.add_argument("--user", default="admin")
    s.add_argument("--password", default=None)
    s.set_defaults(func=cmd_init)

    s = sub.add_parser("serve", help="启动 Web 服务")
    s.add_argument("--host", default=None)
    s.add_argument("--port", type=int, default=None)
    s.add_argument("--debug", action="store_true")
    s.add_argument("--no-scheduler", action="store_true", help="不启动进程内调度线程")
    s.set_defaults(func=cmd_serve)

    s = sub.add_parser("collect", help="执行一次增量采集")
    s.set_defaults(func=cmd_collect)

    s = sub.add_parser("migrate-csv", help="从旧版 CSV 导入")
    s.add_argument("path", nargs="?")
    s.set_defaults(func=cmd_migrate_csv)

    s = sub.add_parser("import-xlsx", help="从官网 xlsx 导入")
    s.add_argument("path")
    s.set_defaults(func=cmd_import_xlsx)

    s = sub.add_parser("fill-prompt", help="补全缺失的 User Prompt")
    s.set_defaults(func=cmd_fill_prompt)

    s = sub.add_parser("export-csv", help="导出 CSV")
    s.add_argument("path", nargs="?")
    s.set_defaults(func=cmd_export_csv)

    s = sub.add_parser("stats", help="存档概况")
    s.set_defaults(func=cmd_stats)

    s = sub.add_parser("import-creds", help="从 VSCode 设置导入 cookie / UA 到数据库")
    s.set_defaults(func=cmd_import_creds)

    s = sub.add_parser("passwd", help="重置 / 创建登录账号")
    s.add_argument("user")
    s.add_argument("password", nargs="?")
    s.set_defaults(func=cmd_passwd)

    s = sub.add_parser("status", help="调度与最近采集状态")
    s.set_defaults(func=cmd_status)

    s = sub.add_parser("vacuum", help="整理数据库（checkpoint + VACUUM）")
    s.set_defaults(func=cmd_vacuum)

    args = ap.parse_args()
    if not getattr(args, "func", None):
        ap.print_help()
        return 0
    config.ensure_dirs()
    return args.func(args) or 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        _p("\n已中断")
        sys.exit(130)