#!/usr/bin/python

import psycopg2
import logging
import asyncpg
import asyncio


"""
=======

Issue 1 -  cursors cannot be used due to "multiple active portals"
cannot use cursors like so:
--
    async with conn.transaction():
        curs = await conn.cursor(
            'INSERT INTO accounts2 (id, balance) VALUES ($1, $2) RETURNING id',
            5,
            1000,
            # "pending",
        )
        my_id = (await curs.fetchrow())["id"]

--
throws
--
  File "asyncpg/protocol/protocol.pyx", line 316, in query
asyncpg.exceptions.FeatureNotSupportedError: unimplemented: multiple active portals not supported
HINT:  You have attempted to use a feature that is not yet implemented.
See: https://go.crdb.dev/issue-v/40195/v20.2
--
seems to be identified here: https://github.com/MagicStack/asyncpg/issues/580

=======
Issue 2 - enums not working with prepared statements
running this:
--
    await conn.execute("INSERT INTO accounts2 (id, balance, mode) VALUES ($1, $2, $3)", 9, 2000, "active", )
--
throws
--
    statement = await self._protocol.prepare(stmt_name, query, timeout)
  File "asyncpg/protocol/protocol.pyx", line 163, in prepare
--
the prepared statement works fine in pgsql.  enums work fine without prepared statements.

=======
Issue 3 - ROW contruction syntax doesn't work in crdb

Running this:
--
conn.execute('UPDATE accounts2 SET (balance) = ROW($1) WHERE id=$2', 2222, 9)
--
on cdrb throws this:
--
asyncpg.exceptions.PostgresSyntaxError: at or near "row": syntax error
DETAIL:  source SQL:
UPDATE accounts2 SET (balance) = ROW($1) WHERE id=$2
                                 ^
HINT:  try \h UPDATE
--
but works just fine on pgsql

"""


def create_schema(conn):
    with conn.cursor() as cur:
        # cur.execute("CREATE TYPE mode2 AS ENUM ('pending', 'active', 'locked')")
        cur.execute(
            "CREATE TABLE IF NOT EXISTS accounts2 (id INT PRIMARY KEY, balance INT, goal int, mode mode2)"
        )
        logging.debug("create_schema: status message: %s", cur.statusmessage)
    conn.commit()


def test1_and_2(conn):
    with conn.cursor() as cur:
        res = cur.execute(
            """INSERT INTO accounts2 (id, balance, mode) VALUES (%s, %s, %s) RETURNING id""",
            (4, 100, "active"),
        )
        my_id = cur.fetchone()[0]
        print(f"{my_id=}")
        assert my_id == 4
        logging.debug("create_accounts(): status message: %s", cur.statusmessage)
    conn.commit()


def test3(conn):
    with conn.cursor() as cur:
        res = cur.execute(
            """UPDATE accounts2 SET (balance) = ROW($1) WHERE id=$2""",
            (111, 4),
        )
        print(f"{res=}")
        assert my_id == 4
        logging.debug("create_accounts(): status message: %s", cur.statusmessage)
    conn.commit()


async def test3_a_bad(conn):
    reslist = await conn.execute(
        """UPDATE accounts2 SET (balance) = ROW($1) WHERE id=$2""",
        2222,
        9,
    )
    print(f"{reslist=}")


async def test3_a_working(conn):
    reslist = await conn.execute(
        """UPDATE accounts2 SET (balance, goal) = ($1, $3) WHERE id=$2""",
        1111,
        7,
        99999,
    )
    print(f"{reslist=}")


async def test1_a_bad(conn):
    async with conn.transaction():
        curs = await conn.cursor(
            """INSERT INTO accounts2 (id, balance) VALUES ($1, $2) RETURNING id""",
            5,
            1000,
            # "pending",
        )
        my_id = (await curs.fetchrow())["id"]
        print(f"{my_id=}")
        assert my_id == 5


async def test1_a_working(conn):
    async with conn.transaction():
        reslist = await conn.fetch(
            """INSERT INTO accounts2 (id, balance) VALUES ($1, $2) RETURNING id""",
            7,
            1000,
            # "pending",
        )
        my_id = reslist[0]["id"]
        print(f"{my_id=}")
        assert my_id == 7


async def test2_a_bad(conn):
    reslist = await conn.execute(
        """INSERT INTO accounts2 (id, balance, mode) VALUES ($1, $2, $3)""",
        9,
        2000,
        "active",
    )


async def test2_a_working(conn):
    reslist = await conn.execute(
        """INSERT INTO accounts2 (id, balance, mode) VALUES (11, 1000, 'active')"""
    )


def main_sync(dsn):
    with psycopg2.connect(dsn) as conn:
        create_schema(conn)
        test1_and_2(conn)
        # test3(conn)


async def main_async(dsn):
    conn = await asyncpg.connect(dsn)

    # await test1_a_working(conn)
    # await test1_a_bad(conn)
    await test2_a_working(conn)
    await test2_a_bad(conn)
    # await test3_a_working(conn)  # needs test1_a_working
    # await test3_a_bad(conn)  # needs test2_a_bad
    await conn.close()


if __name__ == "__main__":
    crdb_dsn = "postgresql://banker:gold@localhost:26257/bank?sslmode=require"
    pgsql_dsn = "postgresql://swimmer:lap@localhost:5432/lake"
    main_sync(crdb_dsn)
    asyncio.run(main_async(crdb_dsn))
    # main_sync(pgsql_dsn)
    # asyncio.run(main_async(pgsql_dsn))
    print("fin")