# If you want best perf, the chdb.connect() API set will satisfy you
from chdb import connect

conn = connect(":memory:")
cursor = conn.cursor()
# Create a table
cursor.execute(
    """
    CREATE TABLE users (
        id Int32,
        name String,
        scores Array(UInt8)
    ) ENGINE = Memory
    """
)

# Insert test data
cursor.execute(
    """
    INSERT INTO users VALUES
    (1, 'Alice', [95, 87, 92]),
    (2, 'Bob', [88, 85, 90]),
    (3, 'Charlie', [91, 89, 94])
    """
)

# Test fetchone
cursor.execute("SELECT * FROM users WHERE id = 1")
row = cursor.fetchone()

import tempfile
import unittest
from chdb import dbapi

with tempfile.TemporaryDirectory() as tmpdirname:
    conn = dbapi.connect(tmpdirname)
    print(conn)
    cur = conn.cursor()
    # cur.execute("CREATE DATABASE IF NOT EXISTS test_db ENGINE = Atomic")
    # cur.execute("USE test_db")
    cur.execute(
        """
    CREATE TABLE rate (
        day Date,
        value Int64
    ) ENGINE = ReplacingMergeTree ORDER BY day"""
    )

    # Insert single value
    cur.execute("INSERT INTO rate VALUES (%s, %s)", ("2021-01-01", 24))
    # Insert multiple values
    cur.executemany(
        "INSERT INTO rate VALUES (%s, %s)",
        [("2021-01-02", 128), ("2021-01-03", 256)],
    )
    # Test executemany outside optimized INSERT/REPLACE path
    cur.executemany(
        "ALTER TABLE rate UPDATE value = %s WHERE day = %s",
        [(72, "2021-01-02"), (96, "2021-01-03")],
    )

    # Test fetchone
    cur.execute("SELECT value FROM rate ORDER BY day DESC LIMIT 2")
    row1 = cur.fetchone()
    self.assertEqual(row1, (96,))
    row2 = cur.fetchone()
    self.assertEqual(row2, (72,))
    row3 = cur.fetchone()
    self.assertIsNone(row3)

    # Test fetchmany
    cur.execute("SELECT value FROM rate ORDER BY day DESC")
    result_set1 = cur.fetchmany(2)
    self.assertEqual(result_set1, ((96,), (72,)))
    result_set2 = cur.fetchmany(1)
    self.assertEqual(result_set2, ((24,),))

    # Test fetchall
    cur.execute("SELECT value FROM rate ORDER BY day DESC")
    rows = cur.fetchall()
    self.assertEqual(rows, ((96,), (72,), (24,)))

    # Clean up
    cur.close()
    conn.close()