import psycopg2
import datetime
from bin_public.config.presets import *
from dateutil import tz
import os


def current_time(type):
    if type == 'ymd':
        return datetime.datetime.now().strftime("%Y-%m-%d")
    if type == 'ymdhms':
        return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")


# hologres 基础函数：查询
def holo_query_func(run_sql, is_query=0):
    conn = psycopg2.connect(host=os.environ['HOST'],
                            port=os.environ['PORT'],
                            dbname=os.environ['DBNAME'],
                            user=os.environ['AK'],
                            password=os.environ['SK'])
    cur = conn.cursor()
    cur.execute(run_sql)
    if is_query:
        data = cur.fetchall()
    cur.close()
    conn.close()
    if is_query:
        return data

def holo_query_account_mapping(invite_code):
    run_sql = f"""
        select end_date, status, mapping_ak
        from s_account_invite_code 
        where invite_code = '{invite_code}'
        order by gmt_modify desc
        limit 1
    """
    data = holo_query_func(run_sql, is_query=1)
    # 数据库中查不到，则返回no_invite_code_msg
    if len(data) == 0:
        status_text = standard_error_msg + no_invite_code_msg
        return status_text, None
    # 数据库中查到，判断是否可用
    if len(data) == 1:
        end_date = data[0][0]
        status = data[0][1]
        mapping_ak = data[0][2]
        if end_date < datetime.datetime.now().strftime("%Y%m%d") or status != '1':
            status_text = standard_error_msg + no_useful_invite_code_msg
            return status_text, None
    return 'Success status: ready', mapping_ak


def key_preprocessing(keyTxt):
    invite_code = keyTxt
    # 这里先用这个逻辑，到时候等实际的邀请码来了就改一下这个函数就行
    if keyTxt.startswith("dteam_"):
        status_display, keyTxt = holo_query_account_mapping(keyTxt)
        yield status_display, keyTxt, invite_code
        return
    else:
        if len(keyTxt) != 51:
            status_display = standard_error_msg + no_apikey_msg
            yield status_display, keyTxt, invite_code
            return
        yield 'Success status: ready', keyTxt, invite_code
        return


def holo_query_insert_chat_message(invite_code, prompt, response, all_token_cnt, history):
    run_sql = f"""
        insert into s_account_chat_message(
            gmt_create
            ,invite_code
            ,prompt
            ,response
            ,all_token_cnt
            ,history
            ,chat_seq
            ,log_timestamp
        )
        select 
             '{datetime.datetime.now().replace(tzinfo=tz.gettz('Asina/Shanghai')).strftime("%Y-%m-%d %H:%M:%S")}' as gmt_create
            ,'{str(invite_code).replace("'", '"')}' as invite_code
            ,'{str(prompt).replace("'", '"')}' as prompt
            ,'{str(response).replace("'", '"')}' as response
            ,'{str(all_token_cnt).replace("'", '"')}' as all_token_cnt
            ,'{str(history).replace("'", '"')}' as history
            ,'{len(history)}' as chat_seq
            ,localtimestamp as log_timestamp
    """
    holo_query_func(run_sql, is_query=0)