"""Utils about db connections"""

import os
from datetime import datetime, timezone

import certifi
import duckdb
import pymysql
import pymysql.cursors

from sliutil.secret import get_secret

SLI_DB_HOST = os.getenv("SLI_DB_HOST")
SLI_DB_PORT = 4000
SLI_DB_USER = os.getenv("SLI_DB_USER")
SLI_DB_PASSWD = get_secret("SLI_DB_PASSWD")
SLI_DB_NAME = "dbaas-sli"

AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
AWS_SESSION_TOKEN = os.getenv("AWS_SESSION_TOKEN")
AWS_REGION = os.getenv("AWS_REGION")


def get_tidb_connection(
    host: str = SLI_DB_HOST,
    port: int = SLI_DB_PORT,
    user: str = SLI_DB_USER,
    password: str = SLI_DB_PASSWD,
    database: str = SLI_DB_NAME,
) -> pymysql.connections.Connection:
    """Get connections to tidb"""
    return pymysql.connect(
        host=host,
        port=port,
        user=user,
        password=password,
        database=database,
        charset="utf8mb4",
        cursorclass=pymysql.cursors.DictCursor,
        ssl_ca=certifi.where(),
    )


def get_duckdb_connection(
    aws_access_key_id: str = AWS_ACCESS_KEY_ID,
    aws_secret_access_key: str = AWS_SECRET_ACCESS_KEY,
    aws_session_token: str = AWS_SESSION_TOKEN,
    region: str = AWS_REGION,
    init_s3: bool = True,
) -> duckdb.DuckDBPyConnection:
    """Get connections to duckdb"""
    start_db = datetime.now(timezone.utc).timestamp()
    duckdb_conn = duckdb.connect(database=":memory:")
    home_directory = "/tmp/duckdb/"
    if not os.path.exists(home_directory):
        os.mkdir(home_directory)

    if init_s3:
        duckdb_conn.execute(
            f"""
                SET home_directory='{home_directory}';
                install httpfs;
                load httpfs;
            """
        )
        complete_db = datetime.now(timezone.utc).timestamp()
        print(f"duckdb_init, take {complete_db - start_db:.2f}s")

    duckdb_conn.execute(
        f"""
            SET s3_access_key_id='{aws_access_key_id}';
            SET s3_secret_access_key='{aws_secret_access_key}';
            SET s3_session_token='{aws_session_token}';
            SET s3_region='{region}';
        """
    )

    return duckdb_conn