import os
import sys
import json
import time
import shutil
import requests

REPORT_DIR = sys.path[0] + "/report"
DATA_FILE = sys.path[0] + "/../data/table_info_with_answer.json"
MAX_QUESTION = 50

def get_input_file(report_dir, config_obj):
    answer_file = report_dir + "/input.json" + formatFileSuffix(config_obj)
    data_file = DATA_FILE + formatFileSuffix(config_obj)
    shutil.copyfile(data_file, answer_file)
    return answer_file

def create_dir_if_not_exists(dir_path):
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)
        pass
    return dir_path

def get_cur_datetime(date_format="%Y%m%d%H%m%S"):
    timestamp = time.time()
    time_array = time.localtime(timestamp)
    return time.strftime(date_format, time_array)

def read_json_file(file_path):
    f = open(file_path, encoding = "utf-8")
    content = f.read()
    f.close()
    return json.loads(content)


def write_json_file(file_path, obj, indent = False):
    f = open(file_path, "w")
    if indent:
        content = f.write(json.dumps(obj, indent=2))
        pass
    else:
        content = f.write(json.dumps(obj))
    f.close()
    return file_path

def run_cmd(cmd):
    return os.system(cmd)

def get_sqlite3_table(file_path, table):
    mydb = sqlite3.connect(file_path)
    cur = mydb.cursor()
    cur.execute("select name from sqlite_master where type='table'")
    return

def formatFileSuffix(config_obj):
    return "__top_p_" + format(config_obj["top_p"], '.6f') 

def gen_import_cmd(file_path, database):
    ret_list = []
    mydb = sqlite3.connect(file_path)
    cur = mydb.cursor()
    cur.execute("select name from sqlite_master where type='table'")
    tables = cur.fetchall()
    table_list = [item[0] for item in tables]
    for table in table_list:
        sql_cmd = "select sql from sqlite_master where name='" + table + "'"
        cur.execute(sql_cmd)
        for row in cur:
            sql = row[0].replace("NUMERIC", "INT(11)").replace("INTEGER", "INT(11)")
            item = {"database": database, "sql": sql, "table": table}
            ret_list.append(item)
            pass
        pass
    return ret_list
    
def http_post_with_auth(url, token, data=None):
    headers = {"authorization": token, "content-type": "application/json"}
    resp = requests.post(url, headers=headers, data=json.dumps(data))
    return json.loads(resp.text)

# https://us-west-2.prod.aws.tidbcloud.com/api/v1/dataapps/sqleditor/orgs/1/projects/1372813089195381283/clusters/1379661944639654096/session
def gen_sql_session(config_obj):
    url = config_obj["url_prefix"] + "/orgs/" + config_obj["org_id"] + "/projects/" + config_obj["project_id"] + "/clusters/" + config_obj["cluster_id"] + "/session"
    resp_obj = http_post_with_auth(url, config_obj["token"], {"database":""})
    if "data" not in resp_obj:
        return ""
    return resp_obj["data"]
def query_data_api(config_obj, sql):
    url = config_obj["url_prefix"] + "/orgs/" + config_obj["org_id"] + "/projects/" + config_obj["project_id"] + "/clusters/" + config_obj["cluster_id"] + "/statements"
    resp_obj = http_post_with_auth(url, config_obj["token"], {"sql":sql, "sessionid": config_obj["session_id"]})
    if "data" not in resp_obj:
        return resp_obj
    return resp_obj["data"]

def evaluate(config_obj, data_obj):
    report_obj = {"summary":{}, "data":[]}
    nogen_list = []
    fail_list = []
    total_num = 0
    succ_num = 0
    no_gen_num = 0
    for table_item in data_obj: 
        # use database
        query_data_api(config_obj, "use " + table_item["database"])
        for index in range(len(table_item["answers"])):
            if index > MAX_QUESTION:
                break
            answer = table_item["answers"][index]
            total_num +=1
            if answer.replace("\n", "").rstrip() == "No SQL Generated":
                no_gen_num +=1
                nogen_list.append({"database": table_item["database"], "question": table_item["questions"][index], "answer": answer})
                continue
            resp = query_data_api(config_obj, answer)
            report_item = {"database": table_item["database"], "sql": answer, "resp": resp, "question": table_item["questions"][index]}
            if "query" in resp and resp['query'] == 'Query OK!':
                succ_num +=1
            else:
                print(report_item)
                fail_list.append(report_item)
            report_obj["data"].append(report_item)
            pass
        pass
    report_obj["summary"] = {"total": total_num, "succ": succ_num, "no_gen": no_gen_num}
    return report_obj, nogen_list, fail_list

if len(sys.argv) != 2:
    print("useage: config_file " + str(len(sys.argv)))
    exit(0)
    pass

print("programing started.")
config_file = sys.argv[1]

config_obj = read_json_file(config_file)
top_p = config_obj.get("top_p", 0.4)
report_dir = create_dir_if_not_exists(REPORT_DIR + "/" + get_cur_datetime("%Y%m%d%H"))
data_file = get_input_file(report_dir, config_obj) 


if data_file == "":
    print("no datafile generated.")
    exit(0)

print(data_file)

data_obj = read_json_file(data_file)

if "session_id" not in config_obj or config_obj["session_id"] == "":
    session_id = gen_sql_session(config_obj)
    config_obj["session_id"] = session_id
    write_json_file(config_file, config_obj, True)


report_obj, nogen_list, fail_list = evaluate(config_obj, data_obj)

write_json_file(report_dir + "/output.json" + formatFileSuffix(config_obj), report_obj, True)
write_json_file(report_dir + "/nogen_list.json" + formatFileSuffix(config_obj), nogen_list, True)
write_json_file(report_dir + "/fail_list.json" + formatFileSuffix(config_obj), fail_list, True)