import json
import io
import sys
import base64
import datetime
import hashlib
import hmac
import requests

method = "POST"
host = "visual.volcengineapi.com"
region = "cn-north-1"
endpoint = "https://visual.volcengineapi.com"
service = "cv"


def sign(key, msg):
    return hmac.new(key, msg.encode("utf-8"), hashlib.sha256).digest()


def getSignatureKey(key, dateStamp, regionName, serviceName):
    kDate = sign(key.encode("utf-8"), dateStamp)
    kRegion = sign(kDate, regionName)
    kService = sign(kRegion, serviceName)
    kSigning = sign(kService, "request")
    return kSigning


def formatQuery(parameters):
    request_parameters_init = ""
    for key in sorted(parameters):
        request_parameters_init += key + "=" + parameters[key] + "&"
    request_parameters = request_parameters_init[:-1]
    return request_parameters


def signV4Request(access_key, secret_key, service, req_query, req_body):
    if access_key is None or secret_key is None:
        print("No access key is available.")
        sys.exit()

    t = datetime.datetime.now(datetime.timezone.utc)
    current_date = t.strftime("%Y%m%dT%H%M%SZ")

    datestamp = t.strftime("%Y%m%d")  # Date w/o time, used in credential scope
    canonical_uri = "/"
    canonical_querystring = req_query
    signed_headers = "content-type;host;x-content-sha256;x-date"
    payload_hash = hashlib.sha256(req_body.encode("utf-8")).hexdigest()
    content_type = "application/json"
    canonical_headers = (
        "content-type:"
        + content_type
        + "\n"
        + "host:"
        + host
        + "\n"
        + "x-content-sha256:"
        + payload_hash
        + "\n"
        + "x-date:"
        + current_date
        + "\n"
    )
    canonical_request = (
        method
        + "\n"
        + canonical_uri
        + "\n"
        + canonical_querystring
        + "\n"
        + canonical_headers
        + "\n"
        + signed_headers
        + "\n"
        + payload_hash
    )

    algorithm = "HMAC-SHA256"
    credential_scope = datestamp + "/" + region + "/" + service + "/" + "request"
    string_to_sign = (
        algorithm
        + "\n"
        + current_date
        + "\n"
        + credential_scope
        + "\n"
        + hashlib.sha256(canonical_request.encode("utf-8")).hexdigest()
    )
    signing_key = getSignatureKey(secret_key, datestamp, region, service)

    signature = hmac.new(
        signing_key, (string_to_sign).encode("utf-8"), hashlib.sha256
    ).hexdigest()

    authorization_header = (
        algorithm
        + " "
        + "Credential="
        + access_key
        + "/"
        + credential_scope
        + ", "
        + "SignedHeaders="
        + signed_headers
        + ", "
        + "Signature="
        + signature
    )
    headers = {
        "X-Date": current_date,
        "Authorization": authorization_header,
        "X-Content-Sha256": payload_hash,
        "Content-Type": content_type,
    }

    # ************* SEND THE REQUEST *************
    request_url = endpoint + "?" + canonical_querystring

    try:
        r = requests.post(request_url, headers=headers, data=req_body)
    except Exception as err:
        print(f"error occurred: {err}")
    else:
        print("\nRESPONSE++++++++++++++++++++++++++++++++++++")
        print(f"Response code: {r.status_code}\n")

    return r


def img2b64(path):
    with open(path, "rb") as f:
        img_data = f.read()
    b64img = base64.b64encode(img_data).decode("utf-8")
    return b64img

# 超分辨率文档 https://www.volcengine.com/docs/6793/1354024
# 请求凭证，从访问控制申请 https://console.volcengine.com/iam/keymanage/
access_key = "AKxxx"
secret_key = "xxx=="

# 请求Query，按照接口文档中填入即可
query_params = {
    "Action": "CVProcess",
    "Version": "2022-08-31",
}
formatted_query = formatQuery(query_params)

body_params = {
    "req_key": "lens_nnsr2_pic_common",
    "binary_data_base64": [img2b64("test/6.jpg")],
    # "image_urls": [
    #     "https://lf3-volc-editor.volccdn.com/obj/volcfe/sop-public/upload_867abb997165dca20f8ee835a01303d7.png"
    # ],
    "model_quality": "MQ",  # ["HQ", "MQ", "LQ" ]
    "result_format": 0,  # 0 png 1 jpeg
}
formatted_body = json.dumps(body_params)

r = signV4Request(access_key, secret_key, service, formatted_query, formatted_body)
rimg = r.json()["data"]["binary_data_base64"]

Image.open(io.BytesIO(base64.b64decode(rimg[0])))