import sqlite3
from typing import Dict, List, Tuple
from o11ydiagnosis.subtasks.slowlogconverter.slowlog import SlowLog
from o11ydiagnosis.subtasks.slowlogconverter.slowstats import SlowStats


class LogsDB:
    @staticmethod
    def try_init_tables(db: sqlite3.Connection):
        db.execute("""
            CREATE TABLE IF NOT EXISTS "keys" (
                "key" TEXT PRIMARY KEY
            );
        """)
        db.execute("""
            CREATE TABLE "slowlog" (
                "INSTANCE" TEXT DEFAULT NULL,
                "Time" INTEGER NOT NULL,
                "Txn_start_ts" INTEGER DEFAULT NULL,
                "User" TEXT DEFAULT NULL,
                "Host" TEXT DEFAULT NULL,
                "Conn_ID" INTEGER DEFAULT NULL,
                "Session_alias" TEXT DEFAULT NULL,
                "Exec_retry_count" INTEGER DEFAULT NULL,
                "Exec_retry_time" REAL DEFAULT NULL,
                "Query_time" REAL DEFAULT NULL,
                "Parse_time" REAL DEFAULT NULL,
                "Compile_time" REAL DEFAULT NULL,
                "Rewrite_time" REAL DEFAULT NULL,
                "Preproc_subqueries" INTEGER DEFAULT NULL,
                "Preproc_subqueries_time" REAL DEFAULT NULL,
                "Optimize_time" REAL DEFAULT NULL,
                "Wait_TS" REAL DEFAULT NULL,
                "Prewrite_time" REAL DEFAULT NULL,
                "Wait_prewrite_binlog_time" REAL DEFAULT NULL,
                "Commit_time" REAL DEFAULT NULL,
                "Get_commit_ts_time" REAL DEFAULT NULL,
                "Commit_backoff_time" REAL DEFAULT NULL,
                "Backoff_types" TEXT DEFAULT NULL,
                "Resolve_lock_time" REAL DEFAULT NULL,
                "Local_latch_wait_time" REAL DEFAULT NULL,
                "Write_keys" INTEGER DEFAULT NULL,
                "Write_size" INTEGER DEFAULT NULL,
                "Prewrite_region" INTEGER DEFAULT NULL,
                "Txn_retry" INTEGER DEFAULT NULL,
                "Cop_time" REAL DEFAULT NULL,
                "Process_time" REAL DEFAULT NULL,
                "Wait_time" REAL DEFAULT NULL,
                "Backoff_time" REAL DEFAULT NULL,
                "LockKeys_time" REAL DEFAULT NULL,
                "Request_count" INTEGER DEFAULT NULL,
                "Total_keys" INTEGER DEFAULT NULL,
                "Process_keys" INTEGER DEFAULT NULL,
                "Rocksdb_delete_skipped_count" INTEGER DEFAULT NULL,
                "Rocksdb_key_skipped_count" INTEGER DEFAULT NULL,
                "Rocksdb_block_cache_hit_count" INTEGER DEFAULT NULL,
                "Rocksdb_block_read_count" INTEGER DEFAULT NULL,
                "Rocksdb_block_read_byte" INTEGER DEFAULT NULL,
                "DB" TEXT DEFAULT NULL,
                "Index_names" TEXT DEFAULT NULL,
                "Is_internal" BOOLEAN DEFAULT NULL,
                "Digest" TEXT DEFAULT NULL,
                "Stats" TEXT DEFAULT NULL,
                "Cop_proc_avg" REAL DEFAULT NULL,
                "Cop_proc_p90" REAL DEFAULT NULL,
                "Cop_proc_max" REAL DEFAULT NULL,
                "Cop_proc_addr" TEXT DEFAULT NULL,
                "Cop_wait_avg" REAL DEFAULT NULL,
                "Cop_wait_p90" REAL DEFAULT NULL,
                "Cop_wait_max" REAL DEFAULT NULL,
                "Cop_wait_addr" TEXT DEFAULT NULL,
                "Mem_max" INTEGER DEFAULT NULL,
                "Disk_max" INTEGER DEFAULT NULL,
                "KV_total" REAL DEFAULT NULL,
                "PD_total" REAL DEFAULT NULL,
                "Backoff_total" REAL DEFAULT NULL,
                "Write_sql_response_total" REAL DEFAULT NULL,
                "Result_rows" INTEGER DEFAULT NULL,
                "Warnings" TEXT DEFAULT NULL,
                "Backoff_Detail" TEXT DEFAULT NULL,
                "Prepared" BOOLEAN DEFAULT NULL,
                "Succ" BOOLEAN DEFAULT NULL,
                "IsExplicitTxn" BOOLEAN DEFAULT NULL,
                "IsWriteCacheTable" BOOLEAN DEFAULT NULL,
                "Plan_from_cache" BOOLEAN DEFAULT NULL,
                "Plan_from_binding" BOOLEAN DEFAULT NULL,
                "Has_more_results" BOOLEAN DEFAULT NULL,
                "Plan" TEXT DEFAULT NULL,
                "Plan_digest" TEXT DEFAULT NULL,
                "Binary_plan" TEXT DEFAULT NULL,
                "Prev_stmt" TEXT DEFAULT NULL,
                "Query" TEXT DEFAULT NULL,
                PRIMARY KEY ("Time", "Conn_ID", "Digest")
            );
        """)

    def save_slowlogs(self, slowlogs: List[SlowLog]):
        pass


class StatsDB:
    @staticmethod
    def try_init_tables(db: sqlite3.Connection):
        db.execute("""
            CREATE TABLE IF NOT EXISTS "keys" (
                "key" TEXT PRIMARY KEY
            );
        """)
        db.execute("""
            CREATE TABLE IF NOT EXISTS "meta" (
                "digest" TEXT PRIMARY KEY,
                "text" TEXT NOT NULL,
                "stmt_type" TEXT NOT NULL,
                "internal" BOOLEAN NOT NULL
            );
        """)
        db.execute("""
            CREATE TABLE IF NOT EXISTS "stats" (
                "digest" TEXT PRIMARY KEY,
                "count" INTEGER NOT NULL,
                "query_time_sum" REAL NOT NULL,
                "query_time_max" REAL NOT NULL,
                "mem_max_sum" REAL NOT NULL,
                "mem_max_max" REAL NOT NULL,
                "disk_max_sum" REAL NOT NULL,
                "disk_max_max" REAL NOT NULL
            );
        """)
        db.execute("""
            CREATE TABLE IF NOT EXISTS "stats_by_db" (
                "db" TEXT,
                "digest" TEXT NOT NULL,
                "count" INTEGER NOT NULL,
                "query_time_sum" REAL NOT NULL,
                "query_time_max" REAL NOT NULL,
                "mem_max_sum" REAL NOT NULL,
                "mem_max_max" REAL NOT NULL,
                "disk_max_sum" REAL NOT NULL,
                "disk_max_max" REAL NOT NULL,
                PRIMARY KEY ("db", "digest")
            );
        """)
        db.execute("""
            CREATE TABLE IF NOT EXISTS "metric_count_per_minute" (
                "ts" INTEGER NOT NULL,
                "digest" TEXT NOT NULL,
                "count" INT NOT NULL,
                PRIMARY KEY ("ts", "digest")
            );
        """)

    @staticmethod
    def query_keys(db: sqlite3.Connection) -> List[str]:
        cursor = db.cursor()
        try:
            cursor.execute("""
                SELECT
                    "key"
                FROM
                    "keys";
            """)
            rows = cursor.fetchall()
            return [row[0] for row in rows]
        finally:
            cursor.close()

    @staticmethod
    def query_meta(db: sqlite3.Connection) -> Dict[str, Tuple[str, bool]]:
        cursor = db.cursor()
        try:
            cursor.execute("""
                SELECT
                    "digest", "text", "internal"
                FROM
                    "meta";
            """)
            rows = cursor.fetchall()
            metas = {}
            for row in rows:
                metas[row[0]] = (row[1], row[2])
            return metas
        finally:
            cursor.close()

    @staticmethod
    def query_stats(db: sqlite3.Connection) -> Dict[str, SlowStats]:
        cursor = db.cursor()
        try:
            cursor.execute("""
                SELECT
                    "digest",
                    "count",
                    "query_time_sum",
                    "query_time_max",
                    "mem_max_sum",
                    "mem_max_max",
                    "disk_max_sum",
                    "disk_max_max"
                FROM
                    "stats";
            """)
            rows = cursor.fetchall()
            all_stats = {}
            for row in rows:
                stats = SlowStats(row[0])
                stats.count = row[1]
                stats.query_time_sum = row[2]
                stats.query_time_max = row[3]
                stats.mem_max_sum = row[4]
                stats.mem_max_max = row[5]
                stats.disk_max_sum = row[6]
                stats.disk_max_max = row[7]
                all_stats[stats.digest] = stats
            return all_stats
        finally:
            cursor.close()

    @staticmethod
    def query_stats_by_db(db: sqlite3.Connection) -> Dict[str, Dict[str, SlowStats]]:
        cursor = db.cursor()
        try:
            cursor.execute("""
                SELECT
                    "db",
                    "digest",
                    "count",
                    "query_time_sum",
                    "query_time_max",
                    "mem_max_sum",
                    "mem_max_max",
                    "disk_max_sum",
                    "disk_max_max"
                FROM
                    "stats_by_db";
            """)
            rows = cursor.fetchall()
            all_stats = {}
            for row in rows:
                stats = SlowStats(row[1])
                stats.count = row[2]
                stats.query_time_sum = row[3]
                stats.query_time_max = row[4]
                stats.mem_max_sum = row[5]
                stats.mem_max_max = row[6]
                stats.disk_max_sum = row[7]
                stats.disk_max_max = row[8]
                if row[0] not in all_stats:
                    all_stats[row[0]] = {}
                all_stats[row[0]][stats.digest] = stats
            return all_stats
        finally:
            cursor.close()

    @staticmethod
    def query_count_per_minute(db: sqlite3.Connection) -> Dict[str, Dict[int, int]]:
        cursor = db.cursor()
        try:
            cursor.execute("""
                SELECT
                    "digest",
                    "ts",
                    "count"
                FROM
                    "metric_count_per_minute"
                ORDER BY
                    "ts" ASC;
            """)
            rows = cursor.fetchall()
            results = {}
            for row in rows:
                digest = row[0]
                ts = row[1]
                count = row[2]
                if digest not in results:
                    results[digest] = {}
                results[digest][ts] = count
            return results
        finally:
            cursor.close()

    @staticmethod
    def save_keys(db: sqlite3.Connection, keys: List[str]):
        for key in keys:
            db.execute(
                """
                INSERT OR REPLACE INTO "keys" (
                    "key"
                ) VALUES (
                    ?
                );
            """,
                (key,),
            )

    @staticmethod
    def save_meta(db: sqlite3.Connection, metas: Dict[str, Tuple[str, bool]]):
        cursor = db.cursor()
        try:
            cursor.execute("""SELECT count(*) FROM pragma_table_info("meta") WHERE name = 'stmt_type';""")
            row = cursor.fetchone()
            if (not row) or (row[0] == 0):
                cursor.execute("""ALTER TABLE "meta" ADD COLUMN "stmt_type" TEXT;""")
        finally:
            cursor.close()
        for digest, (text, internal) in metas.items():
            db.execute(
                """
                INSERT OR REPLACE INTO "meta" (
                    "digest", "text", "internal", "stmt_type"
                ) VALUES (
                    ?, ?, ?, ?
                );
            """,
                (digest, text, internal, get_stmt_type(text)),
            )

    @staticmethod
    def save_stats(db: sqlite3.Connection, stats: Dict[str, SlowStats]):
        for s in stats.values():
            db.execute(
                """
                INSERT OR REPLACE INTO "stats" (
                    "digest",
                    "count",
                    "query_time_sum",
                    "query_time_max",
                    "mem_max_sum",
                    "mem_max_max",
                    "disk_max_sum",
                    "disk_max_max"
                ) VALUES (
                    ?, ?, ?, ?, ?, ?, ?, ?
                );
            """,
                (
                    s.digest,
                    s.count,
                    s.query_time_sum,
                    s.query_time_max,
                    s.mem_max_sum,
                    s.mem_max_max,
                    s.disk_max_sum,
                    s.disk_max_max,
                ),
            )

    @staticmethod
    def save_stats_by_db(db: sqlite3.Connection, stats_by_db: Dict[str, Dict[str, SlowStats]]):
        for dbname, stats in stats_by_db.items():
            for s in stats.values():
                db.execute(
                    """
                    INSERT OR REPLACE INTO "stats_by_db" (
                        "db",
                        "digest",
                        "count",
                        "query_time_sum",
                        "query_time_max",
                        "mem_max_sum",
                        "mem_max_max",
                        "disk_max_sum",
                        "disk_max_max"
                    ) VALUES (
                        ?, ?, ?, ?, ?, ?, ?, ?, ?
                    );
                """,
                    (
                        dbname,
                        s.digest,
                        s.count,
                        s.query_time_sum,
                        s.query_time_max,
                        s.mem_max_sum,
                        s.mem_max_max,
                        s.disk_max_sum,
                        s.disk_max_max,
                    ),
                )

    @staticmethod
    def save_count_per_minute(db: sqlite3.Connection, count_per_minute: Dict[str, Dict[int, int]]):
        for digest, values in count_per_minute.items():
            for ts, count in values.items():
                db.execute(
                    """
                    INSERT OR REPLACE INTO "metric_count_per_minute" (
                        "ts", "digest", "count"
                    ) VALUES (
                        ?, ?, ?
                    );
                """,
                    (ts, digest, count),
                )


def get_stmt_type(sql: str) -> str:
    if not sql:
        return "Others"
    sql = sql.strip().lower()
    if sql.startswith("alter"):
        return "AlterTable"
    if sql.startswith("analyze"):
        return "AnalyzeTable"
    if sql.startswith("begin"):
        return "Begin"
    if sql.startswith("change"):
        return "Change"
    if sql.startswith("insert"):
        return "Insert"
    if sql.startswith("update"):
        return "Update"
    if sql.startswith("commit"):
        return "Commit"
    if sql.startswith("delete"):
        return "Delete"
    if sql.startswith("select"):
        return "Select"
    if sql.startswith("show"):
        return "Show"
    if sql.startswith("set"):
        return "Set"
    return "Others"