services/storage/database_storage.py

from __future__ import annotations

import hashlib
import json
from pathlib import Path
from threading import RLock
from typing import Any, Callable, TypeVar

from sqlalchemy import Column, String, Text, create_engine, Integer, inspect, text
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

from services.storage.base import StorageBackend

Base = declarative_base()
ResultT = TypeVar("ResultT")


class AccountModel(Base):
    """账号数据模型"""
    __tablename__ = "accounts"

    id = Column(Integer, primary_key=True, autoincrement=True)
    access_token = Column(String(2048), unique=True, nullable=False, index=True)
    owner_id = Column(String(255), nullable=True, index=True)
    shared = Column(Integer, nullable=False, default=0, index=True)
    status = Column(String(64), nullable=True, index=True)
    data = Column(Text, nullable=False)  # JSON 格式存储完整账号数据


class AuthKeyModel(Base):
    """鉴权密钥数据模型"""
    __tablename__ = "auth_keys"

    id = Column(Integer, primary_key=True, autoincrement=True)
    key_id = Column(String(255), unique=True, nullable=False, index=True)
    data = Column(Text, nullable=False)


class ShopStateModel(Base):
    __tablename__ = "shop_state"

    key = Column(String(64), primary_key=True)
    data = Column(Text, nullable=False)


class RedeemCodeModel(Base):
    __tablename__ = "redeem_codes"

    id = Column(Integer, primary_key=True, autoincrement=True)
    code_id = Column(String(64), unique=True, nullable=False, index=True)
    code_hash = Column(String(128), unique=True, nullable=False, index=True)
    batch_id = Column(String(64), nullable=True, index=True)
    redeemed_by = Column(String(255), nullable=True, index=True)
    status = Column(String(32), nullable=False, index=True)
    data = Column(Text, nullable=False)


class DatabaseStorageBackend(StorageBackend):
    """数据库存储后端（支持 SQLite、PostgreSQL、MySQL 等）"""

    def __init__(self, database_url: str):
        self.database_url = database_url
        connect_args: dict[str, Any] = {}
        if database_url.startswith("sqlite:///"):
            db_path = database_url.removeprefix("sqlite:///")
            if db_path and db_path != ":memory:":
                Path(db_path).parent.mkdir(parents=True, exist_ok=True)
            connect_args["check_same_thread"] = False
        self.engine = create_engine(
            database_url,
            pool_pre_ping=True,  # 自动检测连接是否有效
            pool_recycle=3600,   # 1小时回收连接
            connect_args=connect_args,
        )
        Base.metadata.create_all(self.engine)
        self.Session = sessionmaker(bind=self.engine)
        self._lock = RLock()
        self._migrate_account_query_columns()
        self._migrate_legacy_redeem_codes()

    def load_accounts(self) -> list[dict[str, Any]]:
        """从数据库加载账号数据"""
        session = self.Session()
        try:
            accounts = []
            for row in session.query(AccountModel).all():
                try:
                    account_data = json.loads(row.data)
                    if isinstance(account_data, dict):
                        accounts.append(account_data)
                except json.JSONDecodeError:
                    continue
            return accounts
        finally:
            session.close()

    @staticmethod
    def _account_owner_id(item: dict[str, Any]) -> str:
        return str(item.get("owner_id") or "").strip()

    @staticmethod
    def _account_shared(item: dict[str, Any]) -> int:
        return 1 if bool(item.get("shared")) else 0

    @staticmethod
    def _account_status(item: dict[str, Any]) -> str:
        return str(item.get("status") or "").strip()

    def _account_from_row(self, row: AccountModel) -> dict[str, Any] | None:
        try:
            account_data = json.loads(row.data)
        except json.JSONDecodeError:
            return None
        return account_data if isinstance(account_data, dict) else None

    def _apply_account_row(self, row: AccountModel, item: dict[str, Any]) -> None:
        row.access_token = str(item.get("access_token") or "").strip()
        row.owner_id = self._account_owner_id(item) or None
        row.shared = self._account_shared(item)
        row.status = self._account_status(item) or None
        row.data = json.dumps(item, ensure_ascii=False)

    def _load_account_rows(self, query: Any) -> list[dict[str, Any]]:
        accounts: list[dict[str, Any]] = []
        for row in query.all():
            account_data = self._account_from_row(row)
            if account_data is not None:
                accounts.append(account_data)
        return accounts

    def load_public_accounts(self) -> list[dict[str, Any]]:
        session = self.Session()
        try:
            return self._load_account_rows(
                session.query(AccountModel).filter(
                    (AccountModel.owner_id.is_(None)) | (AccountModel.owner_id == "")
                )
            )
        finally:
            session.close()

    def load_shared_accounts(self) -> list[dict[str, Any]]:
        session = self.Session()
        try:
            return self._load_account_rows(
                session.query(AccountModel).filter(
                    AccountModel.owner_id.isnot(None),
                    AccountModel.owner_id != "",
                    AccountModel.shared == 1,
                )
            )
        finally:
            session.close()

    def load_accounts_by_owner(self, owner_id: str) -> list[dict[str, Any]]:
        normalized_owner = str(owner_id or "").strip()
        if not normalized_owner:
            return []
        session = self.Session()
        try:
            return self._load_account_rows(session.query(AccountModel).filter_by(owner_id=normalized_owner))
        finally:
            session.close()

    def load_account(self, access_token: str) -> dict[str, Any] | None:
        normalized_token = str(access_token or "").strip()
        if not normalized_token:
            return None
        session = self.Session()
        try:
            row = session.query(AccountModel).filter_by(access_token=normalized_token).first()
            return self._account_from_row(row) if row is not None else None
        finally:
            session.close()

    def upsert_accounts(self, accounts: list[dict[str, Any]]) -> None:
        with self._lock:
            session = self.Session()
            try:
                self._acquire_transaction_lock(session, "accounts")
                for item in accounts:
                    if not isinstance(item, dict):
                        continue
                    access_token = str(item.get("access_token") or "").strip()
                    if not access_token:
                        continue
                    row = session.query(AccountModel).filter_by(access_token=access_token).first()
                    if row is None:
                        row = AccountModel(access_token=access_token, data="{}")
                        session.add(row)
                    self._apply_account_row(row, item)
                session.commit()
            except Exception as e:
                session.rollback()
                raise e
            finally:
                session.close()

    def delete_accounts_by_tokens(
        self,
        tokens: list[str],
        *,
        owner_id: str = "",
        admin: bool = True,
    ) -> None:
        normalized_tokens = [str(token or "").strip() for token in dict.fromkeys(tokens) if str(token or "").strip()]
        if not normalized_tokens:
            return
        normalized_owner = str(owner_id or "").strip()
        with self._lock:
            session = self.Session()
            try:
                self._acquire_transaction_lock(session, "accounts")
                query = session.query(AccountModel).filter(AccountModel.access_token.in_(normalized_tokens))
                if not admin:
                    query = query.filter_by(owner_id=normalized_owner)
                query.delete(synchronize_session=False)
                session.commit()
            except Exception as e:
                session.rollback()
                raise e
            finally:
                session.close()

    def save_accounts(self, accounts: list[dict[str, Any]]) -> None:
        """保存账号数据到数据库"""
        with self._lock:
            session = self.Session()
            try:
                self._acquire_transaction_lock(session, "accounts")
                session.query(AccountModel).delete()
                for item in accounts:
                    if not isinstance(item, dict):
                        continue
                    access_token = str(item.get("access_token") or "").strip()
                    if not access_token:
                        continue
                    row = AccountModel(access_token=access_token, data="{}")
                    self._apply_account_row(row, item)
                    session.add(row)
                session.commit()
            except Exception as e:
                session.rollback()
                raise e
            finally:
                session.close()

    def load_auth_keys(self) -> list[dict[str, Any]]:
        """从数据库加载鉴权密钥数据"""
        return self._load_rows(AuthKeyModel)

    def save_auth_keys(self, auth_keys: list[dict[str, Any]]) -> None:
        """保存鉴权密钥数据到数据库"""
        with self._lock:
            self._save_rows(AuthKeyModel, auth_keys, "id", "key_id")

    def merge_auth_keys(self, changed_items: list[dict[str, Any]], deleted_ids: list[str]) -> None:
        """Apply auth-key changes without rewriting unchanged rows.

        This is safer for shared PostgreSQL deployments where multiple Spaces can
        update different users at the same time.
        """
        with self._lock:
            session = self.Session()
            try:
                self._acquire_transaction_lock(session, "auth_keys")
                for key_id in {str(item or "").strip() for item in deleted_ids if str(item or "").strip()}:
                    session.query(AuthKeyModel).filter_by(key_id=key_id).delete()

                for item in changed_items:
                    if not isinstance(item, dict):
                        continue
                    key_id = str(item.get("id") or "").strip()
                    if not key_id:
                        continue
                    payload = json.dumps(item, ensure_ascii=False)
                    row = session.query(AuthKeyModel).filter_by(key_id=key_id).first()
                    if row is None:
                        session.add(AuthKeyModel(key_id=key_id, data=payload))
                    else:
                        row.data = payload
                session.commit()
            except Exception as e:
                session.rollback()
                raise e
            finally:
                session.close()

    def update_auth_keys_locked(
        self,
        mutator: Callable[[list[dict[str, Any]]], tuple[list[dict[str, Any]], ResultT]],
    ) -> ResultT:
        """Load, mutate, and save auth keys under one database transaction lock."""
        with self._lock:
            session = self.Session()
            try:
                self._acquire_transaction_lock(session, "auth_keys")
                rows = session.query(AuthKeyModel).all()
                items: list[dict[str, Any]] = []
                for row in rows:
                    try:
                        item_data = json.loads(row.data)
                    except json.JSONDecodeError:
                        continue
                    if isinstance(item_data, dict):
                        items.append(item_data)

                next_items, result = mutator(items)
                normalized_items = [item for item in next_items if isinstance(item, dict)]
                next_ids = {str(item.get("id") or "").strip() for item in normalized_items}
                next_ids.discard("")

                for row in rows:
                    if str(row.key_id or "").strip() not in next_ids:
                        session.delete(row)

                existing_by_id = {str(row.key_id or "").strip(): row for row in rows}
                for item in normalized_items:
                    key_id = str(item.get("id") or "").strip()
                    if not key_id:
                        continue
                    payload = json.dumps(item, ensure_ascii=False)
                    row = existing_by_id.get(key_id)
                    if row is None:
                        session.add(AuthKeyModel(key_id=key_id, data=payload))
                    else:
                        row.data = payload
                session.commit()
                return result
            except Exception as e:
                session.rollback()
                raise e
            finally:
                session.close()

    @staticmethod
    def _loads_dict(value: str | None) -> dict[str, Any]:
        try:
            data = json.loads(value or "{}")
        except json.JSONDecodeError:
            return {}
        return data if isinstance(data, dict) else {}

    @staticmethod
    def _compact_shop_state(state: dict[str, Any]) -> dict[str, Any]:
        compact = dict(state if isinstance(state, dict) else {})
        codes = compact.get("codes") if isinstance(compact.get("codes"), list) else []
        compact["codes"] = [
            item
            for item in codes
            if isinstance(item, dict) and not (str(item.get("id") or "").strip() and str(item.get("code_hash") or "").strip())
        ]
        return compact

    @staticmethod
    def _code_identity(item: dict[str, Any]) -> tuple[str, str]:
        return str(item.get("id") or "").strip(), str(item.get("code_hash") or "").strip()

    def _load_shop_state_from_session(self, session: Any, key: str = "default") -> dict[str, Any]:
        row = session.query(ShopStateModel).filter_by(key=str(key or "").strip() or "default").first()
        return self._loads_dict(row.data if row is not None else None)

    def _save_shop_state_to_session(self, session: Any, key: str, state: dict[str, Any]) -> None:
        normalized_key = str(key or "").strip() or "default"
        payload_state = self._compact_shop_state(state) if normalized_key == "default" else state
        payload = json.dumps(payload_state if isinstance(payload_state, dict) else {}, ensure_ascii=False)
        row = session.query(ShopStateModel).filter_by(key=normalized_key).first()
        if row is None:
            session.add(ShopStateModel(key=normalized_key, data=payload))
        else:
            row.data = payload

    def _redeem_code_from_row(self, row: RedeemCodeModel) -> dict[str, Any]:
        item = self._loads_dict(row.data)
        item.setdefault("id", row.code_id)
        item.setdefault("code_hash", row.code_hash)
        item.setdefault("status", row.status)
        if row.batch_id and not item.get("batch_id"):
            item["batch_id"] = row.batch_id
        if row.redeemed_by and not item.get("redeemed_by"):
            item["redeemed_by"] = row.redeemed_by
        return item

    def _load_redeem_codes_from_session(self, session: Any) -> list[dict[str, Any]]:
        items: list[dict[str, Any]] = []
        for row in session.query(RedeemCodeModel).order_by(RedeemCodeModel.id.asc()).all():
            item = self._redeem_code_from_row(row)
            code_id, code_hash = self._code_identity(item)
            if code_id and code_hash:
                items.append(item)
        return items

    def _apply_redeem_code_row(self, row: RedeemCodeModel, item: dict[str, Any]) -> None:
        code_id, code_hash = self._code_identity(item)
        row.code_id = code_id
        row.code_hash = code_hash
        row.batch_id = str(item.get("batch_id") or "").strip() or None
        row.redeemed_by = str(item.get("redeemed_by") or "").strip() or None
        row.status = str(item.get("status") or "active").strip() or "active"
        row.data = json.dumps(item, ensure_ascii=False)

    def _sync_redeem_codes_to_session(self, session: Any, codes: list[dict[str, Any]]) -> None:
        normalized_items: list[dict[str, Any]] = []
        seen_hashes: set[str] = set()
        for item in codes:
            if not isinstance(item, dict):
                continue
            code_id, code_hash = self._code_identity(item)
            if not code_id or not code_hash or code_hash in seen_hashes:
                continue
            seen_hashes.add(code_hash)
            normalized_items.append(dict(item))

        rows = session.query(RedeemCodeModel).all()
        existing_by_hash = {str(row.code_hash or "").strip(): row for row in rows}
        next_hashes = {str(item.get("code_hash") or "").strip() for item in normalized_items}
        for row in rows:
            if str(row.code_hash or "").strip() not in next_hashes:
                session.delete(row)

        for item in normalized_items:
            code_hash = str(item.get("code_hash") or "").strip()
            row = existing_by_hash.get(code_hash)
            if row is None:
                row = RedeemCodeModel(
                    code_id=str(item.get("id") or "").strip(),
                    code_hash=code_hash,
                    status=str(item.get("status") or "active").strip() or "active",
                    data="{}",
                )
                session.add(row)
            self._apply_redeem_code_row(row, item)

    def _migrate_account_query_columns(self) -> None:
        with self._lock:
            session = self.Session()
            try:
                existing_columns = {column["name"] for column in inspect(self.engine).get_columns("accounts")}
                column_statements = {
                    "owner_id": "ALTER TABLE accounts ADD COLUMN owner_id VARCHAR(255)",
                    "shared": "ALTER TABLE accounts ADD COLUMN shared INTEGER NOT NULL DEFAULT 0",
                    "status": "ALTER TABLE accounts ADD COLUMN status VARCHAR(64)",
                }
                for column_name, statement in column_statements.items():
                    if column_name not in existing_columns:
                        try:
                            session.execute(text(statement))
                            session.commit()
                        except Exception as exc:
                            session.rollback()
                            refreshed_columns = {
                                column["name"] for column in inspect(self.engine).get_columns("accounts")
                            }
                            if column_name not in refreshed_columns:
                                print(f"[storage] account query column migration skipped: {exc}")
                                return
            except Exception as exc:
                session.rollback()
                print(f"[storage] account query column migration skipped: {exc}")
                return
            finally:
                session.close()

        with self._lock:
            session = self.Session()
            try:
                for statement in (
                    "CREATE INDEX IF NOT EXISTS ix_accounts_owner_id ON accounts (owner_id)",
                    "CREATE INDEX IF NOT EXISTS ix_accounts_shared ON accounts (shared)",
                    "CREATE INDEX IF NOT EXISTS ix_accounts_status ON accounts (status)",
                ):
                    try:
                        session.execute(text(statement))
                    except Exception as exc:
                        print(f"[storage] account query index migration skipped: {exc}")
                rows = session.query(AccountModel).all()
                for row in rows:
                    item = self._account_from_row(row)
                    if item is not None:
                        self._apply_account_row(row, item)
                session.commit()
            except Exception as exc:
                session.rollback()
                print(f"[storage] account query metadata backfill skipped: {exc}")
            finally:
                session.close()

    def _migrate_legacy_redeem_codes(self) -> None:
        with self._lock:
            session = self.Session()
            try:
                self._acquire_transaction_lock(session, "shop_state")
                self._acquire_transaction_lock(session, "redeem_codes")
                state = self._load_shop_state_from_session(session, "default")
                legacy_codes = state.get("codes") if isinstance(state.get("codes"), list) else []
                migrated = 0
                for item in legacy_codes:
                    if not isinstance(item, dict):
                        continue
                    code_id, code_hash = self._code_identity(item)
                    if not code_id or not code_hash:
                        continue
                    row = session.query(RedeemCodeModel).filter_by(code_hash=code_hash).first()
                    if row is not None:
                        continue
                    row = RedeemCodeModel(code_id=code_id, code_hash=code_hash, status="active", data="{}")
                    self._apply_redeem_code_row(row, dict(item))
                    session.add(row)
                    migrated += 1
                if legacy_codes:
                    self._save_shop_state_to_session(session, "default", state)
                session.commit()
                if migrated:
                    print(f"[storage] Migrated {migrated} redeem codes into indexed table")
            except Exception as e:
                session.rollback()
                print(f"[storage] redeem code migration skipped: {e}")
            finally:
                session.close()

    def load_redeem_codes(self) -> list[dict[str, Any]]:
        session = self.Session()
        try:
            return self._load_redeem_codes_from_session(session)
        finally:
            session.close()

    def load_shop_ledger(self) -> list[dict[str, Any]]:
        session = self.Session()
        try:
            state = self._load_shop_state_from_session(session, "default")
            ledger = state.get("ledger") if isinstance(state.get("ledger"), list) else []
            return [item for item in ledger if isinstance(item, dict)]
        finally:
            session.close()

    def get_redeem_code_by_hash(self, code_hash: str) -> dict[str, Any] | None:
        normalized_hash = str(code_hash or "").strip()
        if not normalized_hash:
            return None
        session = self.Session()
        try:
            row = session.query(RedeemCodeModel).filter_by(code_hash=normalized_hash).first()
            return self._redeem_code_from_row(row) if row is not None else None
        finally:
            session.close()

    def has_redeemed_code_in_batch(self, batch_id: str, user_id: str) -> bool:
        normalized_batch = str(batch_id or "").strip()
        normalized_user = str(user_id or "").strip()
        if not normalized_batch or not normalized_user:
            return False
        session = self.Session()
        try:
            return (
                session.query(RedeemCodeModel)
                .filter_by(batch_id=normalized_batch, redeemed_by=normalized_user, status="redeemed")
                .first()
                is not None
            )
        finally:
            session.close()

    def update_redeem_codes_locked(
        self,
        mutator: Callable[[dict[str, Any]], tuple[dict[str, Any], ResultT]],
    ) -> ResultT:
        """Load, mutate, and save redeem codes plus shop ledger in one transaction."""
        with self._lock:
            session = self.Session()
            try:
                self._acquire_transaction_lock(session, "shop_state")
                self._acquire_transaction_lock(session, "redeem_codes")
                state = self._load_shop_state_from_session(session, "default")
                codes = self._load_redeem_codes_from_session(session)
                legacy_codes = state.get("codes") if isinstance(state.get("codes"), list) else []
                state["codes"] = [item for item in legacy_codes if isinstance(item, dict)] + codes
                next_state, result = mutator(state)
                next_state = next_state if isinstance(next_state, dict) else {}
                if "codes" not in next_state:
                    next_state["codes"] = state.get("codes", [])
                next_codes = next_state.get("codes") if isinstance(next_state.get("codes"), list) else []
                self._sync_redeem_codes_to_session(session, next_codes)
                self._save_shop_state_to_session(session, "default", next_state)
                session.commit()
                return result
            except Exception as e:
                session.rollback()
                raise e
            finally:
                session.close()

    def update_user_and_redeem_code_locked(
        self,
        *,
        code_hash: str,
        user_id: str,
        mutator: Callable[
            [dict[str, Any] | None, dict[str, Any] | None, list[dict[str, Any]], dict[str, Any]],
            tuple[dict[str, Any] | None, dict[str, Any] | None, dict[str, Any], ResultT],
        ],
    ) -> ResultT:
        """Update one user and one redeem code without loading every code."""
        normalized_hash = str(code_hash or "").strip()
        normalized_user = str(user_id or "").strip()
        with self._lock:
            session = self.Session()
            try:
                self._acquire_transaction_lock(session, "shop_state")
                self._acquire_transaction_lock(session, "auth_keys")
                self._acquire_transaction_lock(session, "redeem_codes")

                auth_row = session.query(AuthKeyModel).filter_by(key_id=normalized_user).first()
                auth_item = self._loads_dict(auth_row.data) if auth_row is not None else None

                code_row = session.query(RedeemCodeModel).filter_by(code_hash=normalized_hash).first()
                code_item = self._redeem_code_from_row(code_row) if code_row is not None else None

                batch_items: list[dict[str, Any]] = []
                batch_id = str((code_item or {}).get("batch_id") or "").strip()
                if batch_id and normalized_user:
                    batch_rows = (
                        session.query(RedeemCodeModel)
                        .filter_by(batch_id=batch_id, redeemed_by=normalized_user, status="redeemed")
                        .all()
                    )
                    batch_items = [self._redeem_code_from_row(row) for row in batch_rows]

                state = self._load_shop_state_from_session(session, "default")
                next_auth_item, next_code_item, next_state, result = mutator(auth_item, code_item, batch_items, state)

                if next_auth_item is not None:
                    key_id = str(next_auth_item.get("id") or "").strip()
                    if not key_id:
                        raise ValueError("用户不存在")
                    payload = json.dumps(next_auth_item, ensure_ascii=False)
                    if auth_row is None:
                        session.add(AuthKeyModel(key_id=key_id, data=payload))
                    else:
                        auth_row.data = payload

                if next_code_item is not None:
                    next_code_id, next_code_hash = self._code_identity(next_code_item)
                    if not next_code_id or not next_code_hash:
                        raise ValueError("卡密不存在")
                    if code_row is None:
                        code_row = RedeemCodeModel(
                            code_id=next_code_id,
                            code_hash=next_code_hash,
                            status=str(next_code_item.get("status") or "active").strip() or "active",
                            data="{}",
                        )
                        session.add(code_row)
                    self._apply_redeem_code_row(code_row, next_code_item)

                self._save_shop_state_to_session(session, "default", next_state if isinstance(next_state, dict) else {})
                session.commit()
                return result
            except Exception as e:
                session.rollback()
                raise e
            finally:
                session.close()

    def load_shop_state(self) -> dict[str, Any]:
        return self.load_named_state("default")

    def load_named_state(self, key: str) -> dict[str, Any]:
        session = self.Session()
        try:
            normalized_key = str(key or "").strip() or "default"
            state = self._load_shop_state_from_session(session, normalized_key)
            if normalized_key == "default":
                codes = self._load_redeem_codes_from_session(session)
                if codes:
                    legacy_codes = state.get("codes") if isinstance(state.get("codes"), list) else []
                    state["codes"] = [item for item in legacy_codes if isinstance(item, dict)] + codes
            return state
        finally:
            session.close()

    def save_shop_state(self, state: dict[str, Any]) -> None:
        self.save_named_state("default", state)

    def save_named_state(self, key: str, state: dict[str, Any]) -> None:
        with self._lock:
            session = self.Session()
            try:
                normalized_key = str(key or "").strip() or "default"
                if normalized_key == "default" and isinstance(state.get("codes"), list):
                    self._sync_redeem_codes_to_session(session, state["codes"])
                self._save_shop_state_to_session(session, normalized_key, state)
                session.commit()
            except Exception as e:
                session.rollback()
                raise e
            finally:
                session.close()

    def update_shop_state_locked(self, mutator: Callable[[dict[str, Any]], tuple[dict[str, Any], ResultT]]) -> ResultT:
        """Load, mutate, and save shop state under one database transaction lock."""
        with self._lock:
            session = self.Session()
            try:
                self._acquire_transaction_lock(session, "shop_state")
                self._acquire_transaction_lock(session, "redeem_codes")
                state = self._load_shop_state_from_session(session, "default")
                codes = self._load_redeem_codes_from_session(session)
                if codes:
                    legacy_codes = state.get("codes") if isinstance(state.get("codes"), list) else []
                    state["codes"] = [item for item in legacy_codes if isinstance(item, dict)] + codes

                next_state, result = mutator(state)
                next_state = next_state if isinstance(next_state, dict) else {}
                if "codes" not in next_state:
                    next_state["codes"] = state.get("codes", [])
                next_codes = next_state.get("codes") if isinstance(next_state.get("codes"), list) else []
                self._sync_redeem_codes_to_session(session, next_codes)
                self._save_shop_state_to_session(session, "default", next_state)
                session.commit()
                return result
            except Exception as e:
                session.rollback()
                raise e
            finally:
                session.close()

    def update_user_and_shop_state_locked(
        self,
        mutator: Callable[[list[dict[str, Any]], dict[str, Any]], tuple[list[dict[str, Any]], dict[str, Any], ResultT]],
    ) -> ResultT:
        """Load, mutate, and save auth keys plus shop state in one transaction."""
        with self._lock:
            session = self.Session()
            try:
                self._acquire_transaction_lock(session, "shop_state")
                self._acquire_transaction_lock(session, "auth_keys")
                self._acquire_transaction_lock(session, "redeem_codes")

                auth_rows = session.query(AuthKeyModel).all()
                auth_items: list[dict[str, Any]] = []
                for row in auth_rows:
                    try:
                        item_data = json.loads(row.data)
                    except json.JSONDecodeError:
                        continue
                    if isinstance(item_data, dict):
                        auth_items.append(item_data)

                shop_state = self._load_shop_state_from_session(session, "default")
                codes = self._load_redeem_codes_from_session(session)
                if codes:
                    legacy_codes = shop_state.get("codes") if isinstance(shop_state.get("codes"), list) else []
                    shop_state["codes"] = [item for item in legacy_codes if isinstance(item, dict)] + codes

                next_auth_items, next_shop_state, result = mutator(auth_items, shop_state)
                next_shop_state = next_shop_state if isinstance(next_shop_state, dict) else {}
                if "codes" not in next_shop_state:
                    next_shop_state["codes"] = shop_state.get("codes", [])
                normalized_auth_items = [item for item in next_auth_items if isinstance(item, dict)]
                next_ids = {str(item.get("id") or "").strip() for item in normalized_auth_items}
                next_ids.discard("")

                for row in auth_rows:
                    if str(row.key_id or "").strip() not in next_ids:
                        session.delete(row)

                existing_by_id = {str(row.key_id or "").strip(): row for row in auth_rows}
                for item in normalized_auth_items:
                    key_id = str(item.get("id") or "").strip()
                    if not key_id:
                        continue
                    payload = json.dumps(item, ensure_ascii=False)
                    row = existing_by_id.get(key_id)
                    if row is None:
                        session.add(AuthKeyModel(key_id=key_id, data=payload))
                    else:
                        row.data = payload

                next_codes = next_shop_state.get("codes") if isinstance(next_shop_state.get("codes"), list) else []
                self._sync_redeem_codes_to_session(session, next_codes)
                self._save_shop_state_to_session(session, "default", next_shop_state)

                session.commit()
                return result
            except Exception as e:
                session.rollback()
                raise e
            finally:
                session.close()

    def _load_rows(self, model: type[AccountModel] | type[AuthKeyModel]) -> list[dict[str, Any]]:
        session = self.Session()
        try:
            items = []
            for row in session.query(model).all():
                try:
                    item_data = json.loads(row.data)
                    if isinstance(item_data, dict):
                        items.append(item_data)
                except json.JSONDecodeError:
                    continue
            return items
        finally:
            session.close()

    def _save_rows(
        self,
        model: type[AccountModel] | type[AuthKeyModel],
        items: list[dict[str, Any]],
        source_key: str,
        target_key: str | None = None,
    ) -> None:
        session = self.Session()
        try:
            session.query(model).delete()
            for item in items:
                if not isinstance(item, dict):
                    continue
                key_value = str(item.get(source_key) or "").strip()
                if not key_value:
                    continue
                session.add(
                    model(
                        **{target_key or source_key: key_value},
                        data=json.dumps(item, ensure_ascii=False),
                    )
                )
            session.commit()
        except Exception as e:
            session.rollback()
            raise e
        finally:
            session.close()

    @property
    def _is_postgresql(self) -> bool:
        normalized = self.database_url.lower()
        return normalized.startswith("postgresql://") or normalized.startswith("postgres://")

    @staticmethod
    def _lock_id(name: str) -> int:
        digest = hashlib.sha256(f"chatsam:{name}".encode("utf-8")).digest()
        return int.from_bytes(digest[:8], "big") & 0x7FFF_FFFF_FFFF_FFFF

    def _acquire_transaction_lock(self, session: Any, name: str) -> None:
        if not self._is_postgresql:
            return
        session.execute(text("SELECT pg_advisory_xact_lock(:lock_id)"), {"lock_id": self._lock_id(name)})

    def health_check(self) -> dict[str, Any]:
        """健康检查"""
        try:
            session = self.Session()
            try:
                # 尝试执行简单查询
                session.execute(text("SELECT 1"))
                count = session.query(AccountModel).count()
                auth_key_count = session.query(AuthKeyModel).count()
                shop_state_count = session.query(ShopStateModel).count()
                redeem_code_count = session.query(RedeemCodeModel).count()
                return {
                    "status": "healthy",
                    "backend": "database",
                    "database_url": self._mask_password(self.database_url),
                    "account_count": count,
                    "auth_key_count": auth_key_count,
                    "shop_state_count": shop_state_count,
                    "redeem_code_count": redeem_code_count,
                }
            finally:
                session.close()
        except Exception as e:
            return {
                "status": "unhealthy",
                "backend": "database",
                "error": str(e),
            }

    def get_backend_info(self) -> dict[str, Any]:
        """获取存储后端信息"""
        db_type = "unknown"
        if "sqlite" in self.database_url:
            db_type = "sqlite"
        elif "postgresql" in self.database_url or "postgres" in self.database_url:
            db_type = "postgresql"
        elif "mysql" in self.database_url:
            db_type = "mysql"
        
        return {
            "type": "database",
            "db_type": db_type,
            "description": f"数据库存储 ({db_type})",
            "database_url": self._mask_password(self.database_url),
        }

    @staticmethod
    def _mask_password(url: str) -> str:
        """隐藏数据库连接字符串中的密码"""
        if "://" not in url:
            return url
        try:
            protocol, rest = url.split("://", 1)
            if "@" in rest:
                credentials, host = rest.split("@", 1)
                if ":" in credentials:
                    username, _ = credentials.split(":", 1)
                    return f"{protocol}://{username}:****@{host}"
            return url
        except Exception:
            return url