import os
import time
import hashlib
import base64
import hmac
import logging

import requests
from requests.auth import HTTPDigestAuth

base_url = 'https://data.dev.tidbcloud.com/api/v1beta/app/chat2query-JJBAZFwt/endpoint'
headers = {}
public_key, private_key = "O0T10NO0", "4344f3e3-9638-43f2-a079-ffd2dfa9b6e6"
cluster_id = "10549307547250536346"
db_name = "test"

if os.getenv("ENV") == "prod":
    base_url = "https://data.tidbcloud.com/api/v1beta/app/chat2query-lAxeIOlD/endpoint"
    public_key, private_key = "40RURTH0", "4d8748d9-e1ef-427a-8b54-cb05d82ffa2a"
    cluster_id = "1379661944643734204"
    db_name = "acs_1_year_2015"


def gen_sign(timestamp, secret):
    # Splicing timestamp and secret
    string_to_sign = '{}\n{}'.format(timestamp, secret)
    hmac_code = hmac.new(string_to_sign.encode("utf-8"), digestmod=hashlib.sha256).digest()

    # Perform base64 processing on the result
    sign = base64.b64encode(hmac_code).decode('utf-8')

    return sign


def send_feishu_msg(msg: str):
    url = "https://open.feishu.cn/open-apis/bot/v2/hook/729ea554-3523-4829-9cf3-8dfa90ecec5d"
    token = "HO0vM2bkeizmRdw5GqVrn"
    timestamp = str(int(time.time()))
    sign = gen_sign(timestamp, token)
    payload = {
        "timestamp": timestamp,
        "sign": sign,
        "msg_type": "text",
        "content": {
            "text": msg
        }
    }

    resp = requests.post(url, json=payload)
    logging.info("send feishu msg %s got %s", payload, resp.text)


def get_auth_key():
    return HTTPDigestAuth(public_key, private_key)


def query_job_detail_v2(job_id):
    url = f"{base_url}/v2/jobs/{job_id}"
    resp = requests.get(url, headers=headers, auth=get_auth_key())
    logging.info("get job %s detail got resp %s", job_id, resp.text)

    return resp.json()["result"]


def loopv2(job_id: str):
    for _ in range(10000):
        try:
            result = query_job_detail_v2(job_id)
            if result["status"] in ("done", "failed"):
                break
        except Exception as e:
            logging.exception("failed to query job detail: %s", e)
        finally:
            time.sleep(5)


def get_data_summary_by_cluster(cluster_id: str, dbname: str):
    url = f"{base_url}/v2/dataSummaries"
    params = {
        "cluster_id": cluster_id,
        "database": dbname,
    }
    resp = requests.get(url, params=params, auth=get_auth_key(), headers=headers)
    logging.info("get_data_summary_by_cluster %s got %s", params, resp.text)
    return resp.json()


def create_data_summary_by_cluster(cluster_id: str, dbname: str):
    url = f"{base_url}/v2/dataSummaries"

    payload = {
        "cluster_id": cluster_id,
        "database": dbname,
        "force": True,
    }

    resp = requests.post(url, json=payload, auth=get_auth_key(), headers=headers)
    logging.info("create_data_summary %s, %s got %s: %s", url, payload, resp.text, resp.headers)

    resp_json = resp.json()
    return resp_json["result"]["data_summary_id"], resp_json["result"]["job_id"]


def update_data_summary(data_summary_id: int):
    url = f"{base_url}/v2/dataSummaries/{data_summary_id}"
    payload = {
        "tables": {
            "user": {
                "description": "I_CAN_UPDATE_TABLE_DESCRIPTION",
                "columns": {
                    "username": "I_CAN_UPDATE_COLUMN_DESCRIPTION_TOO",
                }
            }
        }
    }
    resp = requests.put(url, auth=get_auth_key(), headers=headers, json=payload)
    logging.info("update data summary of %s got %s", data_summary_id, resp.text)
    return resp.json()["result"]["job_id"]


def get_data_summary(data_summary_id: int):
    url = f"{base_url}/v2/dataSummaries/{data_summary_id}"
    resp = requests.get(url, auth=get_auth_key(), headers=headers)
    logging.info("get data summary of %s got %s", data_summary_id, resp.text)
    return resp.json()


def chat2data(cluster_id: str, database: str, question: str):
    url = f"{base_url}/v2/chat2data"
    payload = {
        "cluster_id": cluster_id,
        "database": database,
        "raw_question": question
    }
    resp = requests.post(url, json=payload, auth=get_auth_key(), headers=headers)
    logging.info("chat2data with %s got %s", payload, resp.text)
    return resp.json()["result"]["job_id"]


def runtest_once():
    # steps:
    # 1. create data summary
    # 2. loopv2
    # 3. update data summary
    # 4. loopv2
    # 5. chat2data
    # 6. loopv2
    try:
        data_summary_id, job_id = create_data_summary_by_cluster(cluster_id=cluster_id, dbname=db_name)
        get_data_summary_by_cluster(cluster_id=cluster_id, dbname=db_name)

        # 2. loopv2
        loopv2(job_id=job_id)

        for _ in range(3):
            # 5. chat2data
            job_id = chat2data(cluster_id, db_name, question="Count places")

            # 6. loopv2
            loopv2(job_id=job_id)

        get_data_summary_by_cluster(cluster_id=cluster_id, dbname=db_name)
    except Exception as e:
        send_feishu_msg("e2e test failed: %s" % e)


def always_runtest():
    while True:
        runtest_once()
        time.sleep(1)


if __name__ == "__main__":
    # data_summary_id = 481296
    # job_id = chat2data(data_summary_id=data_summary_id, question="count users")
    # job_id = update_data_summary(data_summary_id=data_summary_id)
    # loopv2(job_id)
    # get_data_summary(data_summary_id)
    # always_runtest()
    runtest_once()