tw/db/designer_roster_db.py

# -*- coding: utf-8 -*-
"""
设计师派单数据库（SQLite）

同一设计师在不同店铺对应不同 group_id，派单时从在线设计师中轮询。
企微群「上线」/「下线」通过 update_online(wechat_user_id, is_online) 更新。
"""
import sqlite3
import os
from typing import Optional

_DB_PATH = os.path.join(os.path.dirname(__file__), "designer_roster_db", "roster.db")


def _get_conn() -> sqlite3.Connection:
    os.makedirs(os.path.dirname(_DB_PATH), exist_ok=True)
    conn = sqlite3.connect(_DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn


def init_db():
    with _get_conn() as conn:
        conn.execute("""
            CREATE TABLE IF NOT EXISTS designers (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                wechat_user_id TEXT UNIQUE NOT NULL
            )
        """)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS designer_shops (
                designer_id INTEGER NOT NULL,
                shop_id TEXT NOT NULL,
                group_id TEXT NOT NULL,
                PRIMARY KEY (designer_id, shop_id),
                FOREIGN KEY (designer_id) REFERENCES designers(id)
            )
        """)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS designer_online (
                wechat_user_id TEXT PRIMARY KEY,
                is_online INTEGER NOT NULL DEFAULT 0,
                updated_at TEXT
            )
        """)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS round_robin (
                shop_id TEXT PRIMARY KEY,
                last_index INTEGER NOT NULL DEFAULT 0
            )
        """)
        conn.commit()


init_db()


# ========== 设计师管理 ==========

def add_designer(name: str, wechat_user_id: str) -> int:
    """添加设计师，返回 id"""
    with _get_conn() as conn:
        conn.execute(
            "INSERT OR IGNORE INTO designers (name, wechat_user_id) VALUES (?, ?)",
            (name, wechat_user_id),
        )
        conn.commit()
        row = conn.execute("SELECT id FROM designers WHERE wechat_user_id = ?", (wechat_user_id,)).fetchone()
        return row["id"] if row else 0


def set_designer_shop(designer_id: int, shop_id: str, group_id: str):
    """设置设计师在某店铺的分组 ID（同一设计师不同店铺不同 group_id）"""
    with _get_conn() as conn:
        conn.execute(
            "INSERT OR REPLACE INTO designer_shops (designer_id, shop_id, group_id) VALUES (?, ?, ?)",
            (designer_id, shop_id, group_id),
        )
        conn.commit()


def update_online(wechat_user_id: str, is_online: bool):
    """更新设计师在线状态（企微群「上线」/「下线」解析后调用）"""
    from datetime import datetime
    ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    with _get_conn() as conn:
        conn.execute(
            "INSERT OR REPLACE INTO designer_online (wechat_user_id, is_online, updated_at) VALUES (?, ?, ?)",
            (wechat_user_id, 1 if is_online else 0, ts),
        )
        conn.commit()


# ========== 派单 ==========

def get_transfer_group_for_shop(shop_id: str) -> Optional[str]:
    """
    为店铺轮询派单，返回分组 ID。
    从该店铺的在线设计师中轮询选一个，返回其在该店铺的 group_id。
    无人在线则返回 None。
    """
    with _get_conn() as conn:
        rows = conn.execute("""
            SELECT d.wechat_user_id, ds.group_id
            FROM designer_shops ds
            JOIN designers d ON d.id = ds.designer_id
            JOIN designer_online o ON o.wechat_user_id = d.wechat_user_id AND o.is_online = 1
            WHERE ds.shop_id = ?
        """, (shop_id,)).fetchall()

    if not rows:
        return None

    with _get_conn() as conn:
        rr = conn.execute("SELECT last_index FROM round_robin WHERE shop_id = ?", (shop_id,)).fetchone()
        last = rr["last_index"] if rr else 0
        idx = last % len(rows)
        chosen = rows[idx]
        conn.execute(
            "INSERT OR REPLACE INTO round_robin (shop_id, last_index) VALUES (?, ?)",
            (shop_id, idx + 1),
        )
        conn.commit()

    return chosen["group_id"]


# ========== 查询 ==========

def get_all_wechat_user_ids() -> list:
    """获取所有设计师的 wechat_user_id（用于同步在线状态）"""
    with _get_conn() as conn:
        rows = conn.execute("SELECT wechat_user_id FROM designers").fetchall()
        return [r["wechat_user_id"] for r in rows]


def list_designers():
    """列出所有设计师及其店铺分组"""
    with _get_conn() as conn:
        designers = conn.execute("SELECT id, name, wechat_user_id FROM designers").fetchall()
        result = []
        for d in designers:
            shops = conn.execute(
                "SELECT shop_id, group_id FROM designer_shops WHERE designer_id = ?",
                (d["id"],),
            ).fetchall()
            online = conn.execute(
                "SELECT is_online FROM designer_online WHERE wechat_user_id = ?",
                (d["wechat_user_id"],),
            ).fetchone()
            result.append({
                "id": d["id"],
                "name": d["name"],
                "wechat_user_id": d["wechat_user_id"],
                "shops": {s["shop_id"]: s["group_id"] for s in shops},
                "is_online": bool(online and online["is_online"]),
            })
        return result