import asyncio
import sys
from typing import Tuple, Any
from spider.exec_eval import replace_cur_year
import mysql.connector
from mysql.connector import errorcode
import time

TIMEOUT = 60

class TiDB():
    def __init__(self, db_cfg) -> None:
        self.schema_cache = {}
        self.db_cache = {}
        try:
            print("'host': {}, 'port': {}, 'user': {}, 'ssl_ca': {}".format(db_cfg["host"], db_cfg["port"], db_cfg["user"], db_cfg["ssl_ca"]))

            conn = mysql.connector.connect(**db_cfg)
            cursor = conn.cursor()
        except Exception as e:
            print("Something is wrong when connecting to db:", e)
        else:
            self.conn = conn
            self.cursor = cursor
            self.db_cfg = db_cfg

    def get_default_db(self, db_name: str):
        return db_name
    
    def exec_on_db_with_retry(self, db: str, query: str):
        for delay in range(1, 3):
            try:
                self.cursor.execute(query)
                is_success = True
            except Exception as e:
                is_success, exception = False, e

                conn = mysql.connector.connect(**self.db_cfg)
                cursor = conn.cursor()
                self.conn = conn
                self.cursor = cursor
            finally:
                if is_success is True:
                    return
                else:
                    print('TiDB.exec_on_db_with_retry failed, will retry, db {} query {} is_success {} exception {}'.format(db, query, is_success, exception))

                    time.sleep(delay)
                    continue
        else:
            print('TiDB.exec_on_db_with_retry still failed after retry, db {} query {} is_success {} exception {}'.format(db, query, is_success, exception))
            raise exception

    def get_all_dbs(self, default_db: str):
        if default_db in self.db_cache:
            print('get_all_dbs from cache {}'.format(default_db))
            return self.db_cache[default_db]

        print('get_all_dbs {}'.format(default_db))

        query = "select SCHEMA_NAME from INFORMATION_SCHEMA.SCHEMATA where SCHEMA_NAME like '{}%'".format(default_db)
        self.exec_on_db_with_retry('INFORMATION_SCHEMA', query)
        databases = [str(schema[0]) for schema in self.cursor.fetchall()]
        all_dbs = databases[0:1]
        self.db_cache[default_db] = all_dbs
        return all_dbs

    def get_schema(self, db: str):
        """
        Get database's schema, which is a dict with table name as key
        and list of column names as value
        :param db: database path
        :return: schema dict
        """
        if db in self.schema_cache :
            print('get_schema from cache {}'.format(db))
            return self.schema_cache[db]

        print('get_schema from db {}'.format(db))

        schema = {}
        # fetch table names
        query = "select TABLE_NAME from INFORMATION_SCHEMA.TABLES where TABLE_SCHEMA='{}'".format(db)
        self.exec_on_db_with_retry('INFORMATION_SCHEMA', query)
        tables = [str(table[0]) for table in self.cursor.fetchall()]

        # fetch table info
        for table in tables:
            query = "select COLUMN_NAME from INFORMATION_SCHEMA.COLUMNS where TABLE_SCHEMA='{}' and TABLE_NAME='{}'".format(db, table)
            self.exec_on_db_with_retry('INFORMATION_SCHEMA', query)
            schema[table.lower()] = [str(col[0].lower()) for col in self.cursor.fetchall()]

        self.schema_cache[db] = schema
        return schema

    def exec_on_db(self, db: str, query: str) -> Tuple[str, Any]:
        conn = None
        cursor = None
        query = replace_cur_year(query)
        try:
            conn = mysql.connector.connect(**self.db_cfg)
            cursor = conn.cursor()
            cursor.execute("use {}".format(db))
            cursor.execute("SET SESSION MAX_EXECUTION_TIME=60000")
            cursor.execute(query)
            result = cursor.fetchall()
            flag = "result"
        except Exception as e:
            print('exec_on_db exception {}, query {}'.format(e, query), flush=True)
            flag = "exception"
            result = e
        finally:
            if cursor is not None:
                print("cursor.close()")
                cursor.close()
            if conn is not None:
                print("conn.close()")
                conn.close()
            return flag, result

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()

    parser.add_argument('--db_host', dest='db_host', type=str, help="db host")
    parser.add_argument('--db_port', dest='db_port', type=str, help="db port")
    parser.add_argument('--db_user', dest='db_user', type=str, help="db user")
    parser.add_argument('--db_pwd', dest='db_pwd', type=str, help="db password")

    args = parser.parse_args()

    db_accessor = TiDB({'host': args.db_host, 'port': args.db_port, 'user': args.db_user, 'passwd': args.db_pwd})
    
    db_name = 'INFORMATION_SCHEMA'
    default_db = db_accessor.get_default_db(db_name)
    print(default_db)

    all_dbs = db_accessor.get_all_dbs(default_db)
    print(all_dbs)

    schema = db_accessor.get_schema(default_db)
    print(schema)

    flag, denotation = db_accessor.exec_on_db(default_db, 'select * from SCHEMATA limit 1')
    print(flag)
    print(denotation)

    db_accessor.cursor.close()
    db_accessor.conn.close()