import sqlite3
from typing import List, Dict, Any

INT64_MAX = 2**63 - 1


class SlowStatsDB:
    @staticmethod
    def try_init_tables(db: sqlite3.Connection):
        # db.execute("""
        #     CREATE TABLE IF NOT EXISTS "sql_meta" (
        #         "digest" TEXT,
        #         "sql" TEXT
        #     );
        # """)
        # db.execute("""
        #     CREATE TABLE IF NOT EXISTS "plan_meta" (
        #         "digest" TEXT,
        #         "plan" TEXT
        #     );
        # """)
        db.execute("""
            CREATE TABLE IF NOT EXISTS "slow_stats" (
                "serverless_tenant_id" TEXT,
                "serverless_project_id" TEXT,
                "serverless_cluster_id" TEXT,
                "digest" TEXT,
                "query" TEXT,
                "worker_role" TEXT,
                "is_internal" BOOLEAN,
                "is_tiflash" BOOLEAN,
                "count" INT,
                "query_time_sum" REAL,
                "query_time_max" REAL,
                "query_time_p90" REAL,
                "query_time_p99" REAL,
                "mem_max_sum" REAL,
                "mem_max_max" REAL,
                "disk_max_sum" REAL,
                "disk_max_max" REAL
            );
        """)

    # @staticmethod
    # def query_sql_meta(db: sqlite3.Connection) -> List[Dict[str, Any]]:
    #     return query(db, "sql_meta")

    # @staticmethod
    # def query_plan_meta(db: sqlite3.Connection) -> List[Dict[str, Any]]:
    #     return query(db, "plan_meta")

    @staticmethod
    def query_slow_stats(db: sqlite3.Connection) -> List[Dict[str, Any]]:
        return query(db, "slow_stats")

    # @staticmethod
    # def save_sql_meta(db: sqlite3.Connection, meta: List[Dict[str, Any]]):
    #     save(db, "sql_meta", meta)

    # @staticmethod
    # def save_plan_meta(db: sqlite3.Connection, meta: List[Dict[str, Any]]):
    #     save(db, "plan_meta", meta)

    @staticmethod
    def save_slow_stats(db: sqlite3.Connection, stats: List[Dict[str, Any]]):
        save(db, "slow_stats", stats)


class StatementsDB:
    @staticmethod
    def try_init_tables(db: sqlite3.Connection):
        # db.execute("""
        #     CREATE TABLE IF NOT EXISTS "sql_meta" (
        #         "digest" TEXT,
        #         "sql" TEXT
        #     );
        # """)
        db.execute("""
            CREATE TABLE IF NOT EXISTS "statements" (
                "serverless_tenant_id" TEXT,
                "serverless_project_id" TEXT,
                "serverless_cluster_id" TEXT,
                "digest" TEXT,
                "normalized_sql" TEXT,
                "schema_name" TEXT,
                "stmt_type" TEXT,
                "is_internal" BOOLEAN,
                "exec_count" INT,
                "ru" REAL,
                "rru" REAL,
                "wru" REAL
            );
        """)

    # @staticmethod
    # def query_sql_meta(db: sqlite3.Connection) -> List[Dict[str, Any]]:
    #     return query(db, "sql_meta")

    @staticmethod
    def query_statements(db: sqlite3.Connection) -> List[Dict[str, Any]]:
        return query(db, "statements")

    # @staticmethod
    # def save_sql_meta(db: sqlite3.Connection, meta: List[Dict[str, Any]]):
    #     save(db, "sql_meta", meta)

    @staticmethod
    def save_statements(db: sqlite3.Connection, stmts: List[Dict[str, Any]]):
        save(db, "statements", stmts)


def query(db: sqlite3.Connection, table: str) -> List[Dict[str, Any]]:
    cursor = db.cursor()
    try:
        cursor.execute(f"""SELECT * FROM "{table}" """)
        rows = cursor.fetchall()
        column_names = [desc[0] for desc in cursor.description]
        return [dict(zip(column_names, row)) for row in rows]
    finally:
        cursor.close()


def save(db: sqlite3.Connection, table: str, data: List[Dict[str, Any]]):
    for s in data:
        columns = ", ".join(s.keys())
        placeholders = ", ".join("?" for _ in s)
        sql = f"""INSERT OR REPLACE INTO "{table}" ({columns}) VALUES ({placeholders})"""
        db.execute(sql, tuple([max_int(v) for v in s.values()]))


def max_int(v):
    if isinstance(v, int) and v > INT64_MAX:
        return INT64_MAX
    return v