"""04_upsert.py —— ON CONFLICT UPSERT 演示

演示：
  1) DO NOTHING：忽略冲突
  2) DO UPDATE + EXCLUDED：更新冲突行
  3) RETURNING 结合 UPSERT：一次拿到最新行
  4) 基于部分索引的冲突条件
"""
from __future__ import annotations

from _common import connect, print_table, section


DDL = """
DROP TABLE IF EXISTS ch5_up_users;
CREATE TABLE ch5_up_users (
  id         INT PRIMARY KEY,
  name       TEXT NOT NULL,
  email      TEXT NOT NULL,
  updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
  deleted_at TIMESTAMPTZ
);

-- 部分唯一索引：软删除用户的 email 允许重复
CREATE UNIQUE INDEX ch5_up_users_email_active
  ON ch5_up_users (email) WHERE deleted_at IS NULL;
"""


def demo(cur) -> None:
    section("准备表 ch5_up_users 与部分唯一索引")
    cur.execute(DDL)

    section("1. DO NOTHING：主键冲突则忽略")
    cur.execute(
        "INSERT INTO ch5_up_users(id, name, email) VALUES (1, 'Alice', 'a@x.com')"
        " ON CONFLICT (id) DO NOTHING RETURNING id"
    )
    print(f"  首次插入返回: {cur.fetchall()}")

    cur.execute(
        "INSERT INTO ch5_up_users(id, name, email) VALUES (1, 'Alice-dup', 'x')"
        " ON CONFLICT (id) DO NOTHING RETURNING id"
    )
    print(f"  再次插入冲突返回: {cur.fetchall()}  (空代表被忽略)")

    section("2. DO UPDATE + EXCLUDED：冲突则更新")
    cur.execute(
        """
        INSERT INTO ch5_up_users(id, name, email) VALUES (1, 'Alice-v2', 'a2@x.com')
        ON CONFLICT (id) DO UPDATE
        SET name = EXCLUDED.name,
            email = EXCLUDED.email,
            updated_at = NOW()
        RETURNING id, name, email, updated_at
        """
    )
    print_table(cur.fetchall(), ["id", "name", "email", "updated_at"])

    section("3. 基于部分唯一索引的冲突：软删用户 email 可重复")
    cur.execute(
        "INSERT INTO ch5_up_users(id, name, email, deleted_at) VALUES "
        "(2, 'Bob', 'b@x.com', NOW())"
    )
    cur.execute(
        """
        INSERT INTO ch5_up_users(id, name, email) VALUES (3, 'Bob-new', 'b@x.com')
        ON CONFLICT (email) WHERE deleted_at IS NULL DO UPDATE
        SET name = EXCLUDED.name
        RETURNING id, name, email, deleted_at
        """
    )
    print_table(cur.fetchall(), ["id", "name", "email", "deleted_at"])
    print("  说明：软删除的 Bob 不参与唯一约束，所以新 Bob 被成功 INSERT。")

    section("4. 批量 UPSERT：一次 INSERT N 行")
    rows = [(10 + i, f"user_{i}", f"u{i}@x.com") for i in range(5)]
    cur.executemany(
        "INSERT INTO ch5_up_users(id, name, email) VALUES (%s,%s,%s)"
        " ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name",
        rows,
    )
    cur.execute("SELECT id, name, email FROM ch5_up_users ORDER BY id")
    print_table(cur.fetchall(), ["id", "name", "email"])


def main() -> None:
    with connect() as conn:
        with conn.cursor() as cur:
            demo(cur)
            conn.commit()


if __name__ == "__main__":
    main()