xiaobai-review/backend/features/accounts/repository.py

from __future__ import annotations

import sqlite3
from datetime import datetime, timezone
from typing import Any


class AccountRepositoryMixin:
    """Original SQLite account persistence methods, moved without query changes."""

    def count_users(self) -> int:
        with self.connect() as connection:
            row = connection.execute("SELECT COUNT(*) AS total FROM users").fetchone()
        return int(row["total"] if row else 0)

    def first_user_id(self) -> int:
        with self.connect() as connection:
            row = connection.execute("SELECT MIN(id) AS id FROM users").fetchone()
        return int(row["id"] or 0) if row else 0

    def create_user(
        self,
        username: str,
        password_salt: str,
        password_hash: str,
    ) -> dict[str, Any]:
        now = datetime.now(timezone.utc).isoformat(timespec="seconds")
        try:
            with self.connect() as connection:
                role = "admin" if int(connection.execute("SELECT COUNT(*) FROM users").fetchone()[0]) == 0 else "user"
                cursor = connection.execute(
                    """
                    INSERT INTO users
                        (username, password_salt, password_hash, role, created_at, updated_at)
                    VALUES (?, ?, ?, ?, ?, ?)
                    """,
                    (username, password_salt, password_hash, role, now, now),
                )
                user_id = int(cursor.lastrowid)
        except sqlite3.IntegrityError as exc:
            raise ValueError("该账号名已被使用。") from exc
        return {"id": user_id, "username": username, "role": role, "created_at": now}

    def user_by_username(self, username: str) -> dict[str, Any] | None:
        with self.connect() as connection:
            row = connection.execute(
                """
                SELECT id, username, password_salt, password_hash, role, llm_mode,
                       membership_status, membership_plan, membership_starts_at,
                       membership_expires_at, created_at
                FROM users WHERE username = ? COLLATE NOCASE
                """,
                (username,),
            ).fetchone()
        return dict(row) if row else None

    def user_password(self, user_id: int) -> dict[str, str] | None:
        with self.connect() as connection:
            row = connection.execute(
                "SELECT password_salt, password_hash FROM users WHERE id = ?",
                (user_id,),
            ).fetchone()
        return dict(row) if row else None

    def update_user_password(self, user_id: int, password_salt: str, password_hash: str) -> bool:
        now = datetime.now(timezone.utc).isoformat(timespec="seconds")
        with self.connect() as connection:
            cursor = connection.execute(
                "UPDATE users SET password_salt = ?, password_hash = ?, updated_at = ? WHERE id = ?",
                (password_salt, password_hash, now, user_id),
            )
        return cursor.rowcount > 0

    def delete_user(self, user_id: int) -> bool:
        with self.connect() as connection:
            cursor = connection.execute("DELETE FROM users WHERE id = ?", (user_id,))
            return cursor.rowcount > 0

    def create_session(
        self,
        session_hash: str,
        user_id: int,
        csrf_token: str,
        expires_at: str,
    ) -> None:
        now = datetime.now(timezone.utc).isoformat(timespec="seconds")
        with self.connect() as connection:
            connection.execute("DELETE FROM user_sessions WHERE expires_at <= ?", (now,))
            connection.execute(
                """
                INSERT INTO user_sessions
                    (token_hash, user_id, csrf_token, expires_at, created_at, last_seen_at)
                VALUES (?, ?, ?, ?, ?, ?)
                """,
                (session_hash, user_id, csrf_token, expires_at, now, now),
            )

    def session_user(self, session_hash: str) -> dict[str, Any] | None:
        now = datetime.now(timezone.utc).isoformat(timespec="seconds")
        with self.connect() as connection:
            row = connection.execute(
                """
                SELECT u.id, u.username, u.role, u.llm_mode, u.membership_status,
                       u.membership_plan, u.membership_starts_at, u.membership_expires_at,
                       u.created_at, s.csrf_token, s.expires_at
                FROM user_sessions AS s
                JOIN users AS u ON u.id = s.user_id
                WHERE s.token_hash = ? AND s.expires_at > ?
                """,
                (session_hash, now),
            ).fetchone()
            if row:
                connection.execute(
                    "UPDATE user_sessions SET last_seen_at = ? WHERE token_hash = ?",
                    (now, session_hash),
                )
        return dict(row) if row else None

    def delete_session(self, session_hash: str) -> bool:
        with self.connect() as connection:
            cursor = connection.execute(
                "DELETE FROM user_sessions WHERE token_hash = ?",
                (session_hash,),
            )
            return cursor.rowcount > 0

    def get_user_credentials(self, user_id: int) -> str:
        with self.connect() as connection:
            row = connection.execute(
                "SELECT encrypted_payload FROM user_credentials WHERE user_id = ?",
                (user_id,),
            ).fetchone()
        return str(row["encrypted_payload"]) if row else ""

    def save_user_credentials(self, user_id: int, encrypted_payload: str) -> None:
        now = datetime.now(timezone.utc).isoformat(timespec="seconds")
        with self.connect() as connection:
            connection.execute(
                """
                INSERT INTO user_credentials (user_id, encrypted_payload, updated_at)
                VALUES (?, ?, ?)
                ON CONFLICT(user_id) DO UPDATE SET
                    encrypted_payload = excluded.encrypted_payload,
                    updated_at = excluded.updated_at
                """,
                (user_id, encrypted_payload, now),
            )

    def list_user_credentials(self) -> list[dict[str, Any]]:
        with self.connect() as connection:
            rows = connection.execute(
                "SELECT user_id, encrypted_payload FROM user_credentials ORDER BY user_id"
            ).fetchall()
        return [dict(row) for row in rows]

    def user_access(self, user_id: int) -> dict[str, Any] | None:
        with self.connect() as connection:
            row = connection.execute(
                """
                SELECT id, username, role, llm_mode, membership_status, membership_plan,
                       membership_starts_at, membership_expires_at, created_at
                FROM users WHERE id = ?
                """,
                (user_id,),
            ).fetchone()
        return dict(row) if row else None

    def list_users(self) -> list[dict[str, Any]]:
        with self.connect() as connection:
            rows = connection.execute(
                """
                SELECT id, username, role, llm_mode, membership_status, membership_plan,
                       membership_starts_at, membership_expires_at, created_at
                FROM users ORDER BY id
                """
            ).fetchall()
        return [dict(row) for row in rows]

    def update_user_llm_mode(self, user_id: int, mode: str) -> None:
        now = datetime.now(timezone.utc).isoformat(timespec="seconds")
        with self.connect() as connection:
            connection.execute(
                "UPDATE users SET llm_mode = ?, updated_at = ? WHERE id = ?",
                (mode, now, user_id),
            )

    def update_membership(
        self,
        user_id: int,
        status: str,
        plan: str,
        starts_at: str | None,
        expires_at: str | None,
    ) -> bool:
        now = datetime.now(timezone.utc).isoformat(timespec="seconds")
        with self.connect() as connection:
            cursor = connection.execute(
                """
                UPDATE users
                SET membership_status = ?, membership_plan = ?,
                    membership_starts_at = ?, membership_expires_at = ?, updated_at = ?
                WHERE id = ?
                """,
                (status, plan, starts_at, expires_at, now, user_id),
            )
        return cursor.rowcount > 0

    def get_user_birth_profile(self, user_id: int) -> str:
        with self.connect() as connection:
            row = connection.execute(
                "SELECT encrypted_payload FROM user_birth_profiles WHERE user_id = ?",
                (user_id,),
            ).fetchone()
        return str(row["encrypted_payload"]) if row else ""

    def save_user_birth_profile(self, user_id: int, encrypted_payload: str) -> None:
        now = datetime.now(timezone.utc).isoformat(timespec="seconds")
        with self.connect() as connection:
            connection.execute(
                """
                INSERT INTO user_birth_profiles (user_id, encrypted_payload, updated_at)
                VALUES (?, ?, ?)
                ON CONFLICT(user_id) DO UPDATE SET
                    encrypted_payload = excluded.encrypted_payload,
                    updated_at = excluded.updated_at
                """,
                (user_id, encrypted_payload, now),
            )

    def delete_user_birth_profile(self, user_id: int) -> bool:
        with self.connect() as connection:
            cursor = connection.execute(
                "DELETE FROM user_birth_profiles WHERE user_id = ?",
                (user_id,),
            )
            return cursor.rowcount > 0

    def cleanup_expired_switch_grants(self, now: str) -> int:
        with self.connect() as connection:
            cursor = connection.execute(
                "DELETE FROM account_switch_grants WHERE expires_at <= ?",
                (now,),
            )
        return int(cursor.rowcount)

    def list_switch_grants(self, device_hash: str, now: str) -> list[dict[str, Any]]:
        with self.connect() as connection:
            rows = connection.execute(
                """
                SELECT u.id, u.username, u.role, u.llm_mode, u.membership_status,
                       u.membership_plan, u.membership_starts_at, u.membership_expires_at,
                       u.created_at, g.last_used_at, g.granted_at, g.expires_at
                FROM account_switch_grants AS g
                JOIN users AS u ON u.id = g.user_id
                WHERE g.device_hash = ? AND g.expires_at > ?
                ORDER BY g.last_used_at DESC, g.id DESC
                """,
                (device_hash, now),
            ).fetchall()
        return [dict(row) for row in rows]

    def get_switch_grant(self, device_hash: str, user_id: int) -> dict[str, Any] | None:
        with self.connect() as connection:
            row = connection.execute(
                """
                SELECT device_hash, user_id, granted_at, last_used_at, expires_at
                FROM account_switch_grants
                WHERE device_hash = ? AND user_id = ?
                """,
                (device_hash, user_id),
            ).fetchone()
        return dict(row) if row else None

    def upsert_switch_grant(
        self,
        device_hash: str,
        user_id: int,
        granted_at: str,
        last_used_at: str,
        expires_at: str,
    ) -> None:
        with self.connect() as connection:
            connection.execute(
                """
                INSERT INTO account_switch_grants
                    (device_hash, user_id, granted_at, last_used_at, expires_at)
                VALUES (?, ?, ?, ?, ?)
                ON CONFLICT(device_hash, user_id) DO UPDATE SET
                    granted_at = excluded.granted_at,
                    last_used_at = excluded.last_used_at,
                    expires_at = excluded.expires_at
                """,
                (device_hash, user_id, granted_at, last_used_at, expires_at),
            )

    def prune_switch_grants(self, device_hash: str, keep: int) -> int:
        with self.connect() as connection:
            rows = connection.execute(
                """
                SELECT id FROM account_switch_grants
                WHERE device_hash = ?
                ORDER BY last_used_at DESC, id DESC
                """,
                (device_hash,),
            ).fetchall()
            extra = [int(row["id"]) for row in rows[keep:]]
            if not extra:
                return 0
            connection.execute(
                f"DELETE FROM account_switch_grants WHERE id IN ({','.join('?' * len(extra))})",
                extra,
            )
        return len(extra)

    def delete_switch_grant(self, device_hash: str, user_id: int) -> bool:
        with self.connect() as connection:
            cursor = connection.execute(
                "DELETE FROM account_switch_grants WHERE device_hash = ? AND user_id = ?",
                (device_hash, user_id),
            )
        return cursor.rowcount > 0

    def delete_switch_grants_for_user(self, user_id: int) -> int:
        with self.connect() as connection:
            cursor = connection.execute(
                "DELETE FROM account_switch_grants WHERE user_id = ?",
                (user_id,),
            )
        return int(cursor.rowcount)