import time

import pandas as pd

from jcutil.drivers import db


def quick_update(schema, target: str, set_cols, where_cols, data, /, mode="update"):
    """利用pandas快速生成的临时表来批量更新数据

    **使用RPLACE模式需要注意，如果传入的数据字段少于原有字段，则原表中的这些字段将被至为NULL**

    Args:
        schema (str): Database bind alias name
        target (str): update target table name
        set_cols (str, list[str]): Will be updated columns. 'id,name,somecol' or ['id','name','somecol']
        where_cols (str, list[str]): Columns that be used in where. 'id,name,somecol' or ['id', 'name', 'somecol']
        data (Iterable, DataFrame): a data that can pass to pd.DataFrame
        mode (str, optional): sql mode: ['update', 'replace']. Defaults to 'update'.

    Exceptions:
        ValueError: some column in set_cols or where_cols which is not in target table.
    """
    temp_name = f"tmp_{int(time.time())}"
    if isinstance(where_cols, (str, bytes)):
        where_cols = where_cols.split(",")
    if isinstance(set_cols, (str, bytes)):
        set_cols = set_cols.split(",")
    target_columns = set(
        pd.read_sql(f"describe {target}", db.conn(schema)).Field.values
    )
    df = pd.DataFrame(data)
    if diff := (set(set_cols) - target_columns):
        raise ValueError(f"set_cols: {diff} columns not in target table.")
    if diff := (set(where_cols) - target_columns):
        raise ValueError(f"where_cols: {diff} columns not in target table.")
    cols = [col for col in df.columns if col in target_columns]

    if df.empty or not cols:
        return
    if mode == "update":
        sql = "UPDATE {target} t, {source} s SET {sets} WHERE {where}".format(
            target=target,
            source=temp_name,
            sets=",".join([f"t.{c} = s.{c}" for c in set_cols]),
            where=" AND ".join([f"t.{c} = s.{c}" for c in where_cols]),
        )
    else:
        sql = "REPLACE INTO {target} ({sets}) SELECT {sets} FROM {source}".format(
            target=target, source=temp_name, sets=",".join(cols)
        )
    with db.conn(schema).begin() as session:
        session.execute(f"create table {temp_name} like {target}")
        df[cols].to_sql(temp_name, session, index=False, if_exists="append")
        session.execute(sql)
        session.execute(f"drop table {temp_name}")
        session.commit()
