import unittest
import os
import datetime

# Only import if available (for integration tests)
try:
    import duckdb
    from deltalake import DeltaTable
    import pyarrow.compute as pc
    DELTALAKE_AVAILABLE = True
except ImportError:
    DELTALAKE_AVAILABLE = False

S3_PATH_SLOGS = "s3://o11y-prod-shared-ap-east-1/datalake/1372813089209238000/25d59432-9a9e-47cc-880b-d8ddd39458b7/slowlogs"
S3_PATH_QUERIES = "s3://o11y-prod-shared-ap-east-1/datalake/1372813089209238000/25d59432-9a9e-47cc-880b-d8ddd39458b7/queries"

storage_options = {
    "AWS_REGION": "ap-east-1"
}

ZORDER_COLUMNS = ["time", "digest"]
TARGET_PARTITION_DATE = "20250502"


class TestQueryDelta(unittest.TestCase):
    @unittest.skipUnless(DELTALAKE_AVAILABLE, "deltalake not available")
    @unittest.skipUnless(os.environ.get("AWS_ACCESS_KEY_ID"), "AWS credentials not set")
    def test_query_delta(self):
        start_time = datetime.datetime.now()
        print(f"start_time {start_time}")

        dt_slowlogs = DeltaTable(S3_PATH_SLOGS, storage_options=storage_options)
        table_time = datetime.datetime.now()
        print(f"table_time {table_time}")

        dataset = dt_slowlogs.to_pyarrow_dataset()
        pyarrow_time = datetime.datetime.now()
        print(f"pyarrow_time {pyarrow_time}")

        # filter date == '20250523'
        filtered = dataset.to_table(filter=pc.field("date") == "20250523")
        sorted_filtered = filtered.sort_by([("time", "descending")])
        limited = sorted_filtered.slice(0, 100)

        # select columns
        result = limited.select(["digest", "time", "date", "query_time", "max_mem"])
        finish = datetime.datetime.now()
        print(f"finish {finish}")

        # 打印或处理结果
        print(result.to_pandas())

    @unittest.skipUnless(DELTALAKE_AVAILABLE, "deltalake not available")
    @unittest.skipUnless(os.environ.get("AWS_ACCESS_KEY_ID"), "AWS credentials not set")
    def test_query(self):
        try:
            start_time = datetime.datetime.now()

            dt_slowlogs = DeltaTable(S3_PATH_SLOGS, storage_options=storage_options)
            dt_queries = DeltaTable(S3_PATH_QUERIES, storage_options=storage_options)

            con_time = datetime.datetime.now()
            print(f"deltatable time: {(con_time - start_time).total_seconds():.4f} seconds")

            # con = duckdb.connect(database='my_database.duckdb', read_only=False)
            con = duckdb.connect(database=':memory:', read_only=False)
            con.execute("SET memory_limit='8GB';")

            register_time = datetime.datetime.now()
            print(f"connect time: {(register_time - con_time).total_seconds():.4f} seconds")

            con.register("slowlogs", dt_slowlogs.to_pyarrow_dataset())
            con.register("queries", dt_queries.to_pyarrow_dataset())
            end_time = datetime.datetime.now()
            print(f"register time: {(end_time - register_time).total_seconds():.4f} seconds")
            execute(con)
            execute(con)
            execute(con)

            con.close()
        except Exception as e:
            print(f"query failed: {e}")

def execute(con):
    start_time = datetime.datetime.now()
    sql_query = """
    SELECT sl.time, sl.digest, sl.query_time, sl.mem_max
    FROM slowlogs sl
    WHERE sl.date == '20250502'
    order by sl.query_time desc
    LIMIT 100
    """
    # sql_query = """
    # SELECT sl.time, sl.query_time, sl.mem_max, sl.process_keys, sl.total_keys
    # FROM slowlogs sl
    # WHERE sl.date == '20250502'
    # order by time desc
    # """

    result_df = con.execute(sql_query).fetchdf()
    end_time = datetime.datetime.now()

    query_duration = end_time - start_time
    print(f"query time: {query_duration.total_seconds():.4f} seconds")
    #
    print("query result:")
    print(result_df)

if __name__ == "__main__":
    unittest.main()