import ibm_db
import ibm_db_dbi
import pandas as pd

from testcontainers.db2 import Db2Container


def basic_example():
    with Db2Container() as db2:
        # Get connection parameters
        host = db2.get_container_host_ip()
        port = db2.get_exposed_port(db2.port)
        database = db2.database
        username = db2.username
        password = db2.password

        # Create connection string
        conn_str = f"DATABASE={database};HOSTNAME={host};PORT={port};PROTOCOL=TCPIP;UID={username};PWD={password}"

        # Connect to DB2
        conn = ibm_db.connect(conn_str, "", "")
        print("Connected to DB2")

        # Create a test table
        create_table_sql = """
        CREATE TABLE test_table (
            id INTEGER NOT NULL PRIMARY KEY,
            name VARCHAR(50),
            value DECIMAL(10,2),
            created_at TIMESTAMP DEFAULT CURRENT TIMESTAMP
        )
        """

        try:
            ibm_db.exec_immediate(conn, create_table_sql)
            print("Created test table")
        except Exception as e:
            print(f"Table might already exist: {e}")

        # Insert test data
        test_data = [(1, "test1", 100.0), (2, "test2", 200.0), (3, "test3", 300.0)]

        insert_sql = "INSERT INTO test_table (id, name, value) VALUES (?, ?, ?)"
        stmt = ibm_db.prepare(conn, insert_sql)

        for row in test_data:
            ibm_db.execute(stmt, row)
        print("Inserted test data")

        # Query data using ibm_db_dbi
        conn_dbi = ibm_db_dbi.Connection(conn)
        cursor = conn_dbi.cursor()

        cursor.execute("SELECT * FROM test_table ORDER BY id")
        rows = cursor.fetchall()

        print("\nQuery results:")
        for row in rows:
            print(f"ID: {row[0]}, Name: {row[1]}, Value: {row[2]}, Created: {row[3]}")

        # Execute a more complex query
        cursor.execute("""
            SELECT
                name,
                AVG(value) as avg_value,
                COUNT(*) as count,
                MIN(created_at) as first_created,
                MAX(created_at) as last_created
            FROM test_table
            GROUP BY name
            ORDER BY avg_value DESC
        """)

        print("\nAggregation results:")
        for row in cursor.fetchall():
            print(f"Name: {row[0]}, Avg: {row[1]:.2f}, Count: {row[2]}, First: {row[3]}, Last: {row[4]}")

        # Convert to pandas DataFrame
        df = pd.read_sql("SELECT * FROM test_table ORDER BY id", conn_dbi)
        print("\nDataFrame:")
        print(df)

        # Clean up
        cursor.close()
        ibm_db.close(conn)


if __name__ == "__main__":
    basic_example()