"""连接与读写 IO。PG 用 psycopg(3)，MySQL 用 pymysql。桥不 import 基座/平台代码。"""
from contextlib import contextmanager

import pandas as pd
import psycopg
import pymysql

import config


@contextmanager
def akg_pg_conn():
    c = config.akg_pg()
    conn = psycopg.connect(host=c.host, port=c.port, user=c.user,
                           password=c.password, dbname=c.db)
    try:
        yield conn
    finally:
        conn.close()


def _mysql(c: "config.Conn"):
    return pymysql.connect(host=c.host, port=c.port, user=c.user, password=c.password,
                           database=c.db, charset="utf8mb4", read_timeout=180)


@contextmanager
def _mysql_cm(which: str):
    c = {"heat": config.heat_mysql, "factor": config.factor_mysql,
         "price": config.price_mysql}[which]()
    conn = _mysql(c)
    try:
        yield conn
    finally:
        conn.close()


def read_pg(sql: str, params=None) -> pd.DataFrame:
    with akg_pg_conn() as conn:
        return pd.read_sql(sql, conn, params=params)


def read_mysql(which: str, sql: str, params=None) -> pd.DataFrame:
    with _mysql_cm(which) as conn:
        return pd.read_sql(sql, conn, params=params)


@contextmanager
def factor_conn():
    """写因子库用（需要游标提交）。"""
    with _mysql_cm("factor") as conn:
        yield conn