"""第 4-2 课示例：使用事务和 Repository 完成安全转账。"""

from decimal import Decimal
from pathlib import Path
import tomllib

import psycopg
from psycopg import Connection
from psycopg.rows import dict_row


CONFIG_PATH = Path(__file__).with_name("config.toml")


class TransferError(Exception):
    """表示转账过程中可以预期的业务失败。"""


def load_database_config(config_path: Path) -> dict[str, str | int]:
    """读取第二课本地 TOML 数据库配置。"""
    if not config_path.exists():
        raise RuntimeError(
            "未找到 config.toml，请复制 config.example.toml 并填写练习数据库配置。"
        )

    with config_path.open("rb") as config_file:
        config_data = tomllib.load(config_file)

    database_config = config_data.get("postgresql")
    if not isinstance(database_config, dict):
        raise RuntimeError("config.toml 缺少 [postgresql] 配置节。")

    return database_config


class AccountRepository:
    """封装账户表 SQL，但不自行提交或回滚事务。"""

    def __init__(self, connection: Connection) -> None:
        self.connection = connection

    def create_table(self) -> None:
        """创建本课专用表；表已存在时保持不变。"""
        self.connection.execute(
            """
            CREATE TABLE IF NOT EXISTS course_bank_account (
                account_no VARCHAR(30) PRIMARY KEY,
                owner_name VARCHAR(50) NOT NULL,
                balance NUMERIC(12, 2) NOT NULL CHECK (balance >= 0)
            )
            """
        )

    def reset_course_accounts(self) -> None:
        """只清理 COURSE- 前缀的课程数据，避免影响其他记录。"""
        self.connection.execute(
            "DELETE FROM course_bank_account WHERE account_no LIKE %s",
            ("COURSE-%",),
        )

    def add_accounts(self, accounts: list[tuple[str, str, Decimal]]) -> None:
        """使用 executemany() 批量新增课程账户。"""
        # executemany() 是 Cursor 的方法，因此显式创建并关闭游标。
        with self.connection.cursor() as cursor:
            cursor.executemany(
                """
                INSERT INTO course_bank_account (account_no, owner_name, balance)
                VALUES (%s, %s, %s)
                """,
                accounts,
            )

    def get_balance_for_update(self, account_no: str) -> Decimal:
        """查询并锁定账户，防止并发事务同时修改同一余额。"""
        result = self.connection.execute(
            """
            SELECT balance
            FROM course_bank_account
            WHERE account_no = %s
            FOR UPDATE
            """,
            (account_no,),
        ).fetchone()

        if result is None:
            raise TransferError(f"账户不存在：{account_no}")

        return result[0]

    def change_balance(self, account_no: str, amount: Decimal) -> None:
        """使用数据库加法更新余额，并检查目标账户是否存在。"""
        cursor = self.connection.execute(
            """
            UPDATE course_bank_account
            SET balance = balance + %s
            WHERE account_no = %s
            """,
            (amount, account_no),
        )
        if cursor.rowcount != 1:
            raise TransferError(f"账户不存在：{account_no}")

    def find_course_accounts(self) -> list[dict[str, object]]:
        """按账号查询课程账户，并以字典行返回。"""
        cursor = self.connection.cursor(row_factory=dict_row)
        try:
            cursor.execute(
                """
                SELECT account_no, owner_name, balance
                FROM course_bank_account
                WHERE account_no LIKE %s
                ORDER BY account_no
                """,
                ("COURSE-%",),
            )
            return list(cursor.fetchall())
        finally:
            cursor.close()


class TransferService:
    """组织转账业务规则；事务由调用它的连接上下文统一管理。"""

    def __init__(self, repository: AccountRepository) -> None:
        self.repository = repository

    def transfer(self, source_no: str, target_no: str, amount: Decimal) -> None:
        """在同一事务中完成扣款与入账。"""
        if amount <= 0:
            raise TransferError("转账金额必须大于 0。")

        source_balance = self.repository.get_balance_for_update(source_no)
        self.repository.get_balance_for_update(target_no)

        if source_balance < amount:
            raise TransferError("账户余额不足。")

        self.repository.change_balance(source_no, -amount)
        self.repository.change_balance(target_no, amount)


def print_accounts(title: str, accounts: list[dict[str, object]]) -> None:
    """输出当前课程账户余额。"""
    print(title)
    for account in accounts:
        print(
            f"{account['account_no']}｜{account['owner_name']}｜"
            f"余额：{account['balance']}"
        )


def prepare_data(database_config: dict[str, str | int]) -> None:
    """创建专用表并重置本课固定数据。"""
    with psycopg.connect(**database_config) as connection:
        repository = AccountRepository(connection)
        repository.create_table()
        repository.reset_course_accounts()
        repository.add_accounts(
            [
                ("COURSE-A001", "小明", Decimal("1000.00")),
                ("COURSE-A002", "小红", Decimal("500.00")),
            ]
        )


def run_successful_transfer(database_config: dict[str, str | int]) -> None:
    """演示正常离开连接上下文时自动提交事务。"""
    with psycopg.connect(**database_config) as connection:
        service = TransferService(AccountRepository(connection))
        service.transfer("COURSE-A001", "COURSE-A002", Decimal("200.00"))


def run_failed_transfer(database_config: dict[str, str | int]) -> None:
    """演示异常离开连接上下文时自动回滚整个事务。"""
    try:
        with psycopg.connect(**database_config) as connection:
            repository = AccountRepository(connection)

            # 先执行一条成功更新，再主动触发业务异常。
            # 外层 with 会回滚，因此这 50 元扣款不会保留下来。
            repository.change_balance("COURSE-A001", Decimal("-50.00"))
            raise TransferError("模拟第二步失败，验证前一步更新会被回滚。")
    except TransferError as error:
        print(f"失败事务已回滚：{error}")


def query_accounts(
    database_config: dict[str, str | int],
) -> list[dict[str, object]]:
    """使用独立连接回查已经提交的数据。"""
    with psycopg.connect(**database_config) as connection:
        return AccountRepository(connection).find_course_accounts()


def main() -> None:
    """依次演示初始化、提交、回滚和回查。"""
    try:
        database_config = load_database_config(CONFIG_PATH)
        prepare_data(database_config)
        print_accounts("初始余额：", query_accounts(database_config))

        run_successful_transfer(database_config)
        print_accounts("成功转账 200 元后：", query_accounts(database_config))

        run_failed_transfer(database_config)
        print_accounts("失败事务回滚后：", query_accounts(database_config))
    except (OSError, tomllib.TOMLDecodeError, RuntimeError) as error:
        print(f"配置读取失败：{error}")
    except psycopg.Error as error:
        print(f"数据库访问失败：{error}")


if __name__ == "__main__":
    main()