# coding=utf-8
import math
import os
import time

import ks3.multipart
from ks3 import utils
from ks3.connection import Connection

ak = os.getenv('KS3_IAM_ACCESS_KEY_ID', '<YOUR_ACCESS_KEY>')
sk = os.getenv('KS3_IAM_ACCESS_KEY_SECRET', '<YOUR_SECRET_KEY>')
bucket_name = os.getenv('KS3_BUCKET', '<KS3_BUCKET>')
endpoint = os.getenv('KS3_ENDPOINT', 'ks3-cn-beijing-internal.ksyuncs.com')

key_name = 'test-file'
file_path = './test-file'

# 默认开启自动校验crc
# 关闭需要 conn.enable_crc = False 或 Connection(ak, sk, host=endpoint, enable_crc=False)
conn = Connection(ak, sk, host=endpoint)
bucket = conn.get_bucket(bucket_name)

# 下载对象后，计算服务端返回的crc64值，与本地读取计算的crc64值进行对比
def test_downloadObject_checkingCrc64():
    k = bucket.get_key(key_name)

    # 下载成功之后，会生成server_crc
    # 读取对象数据之后，会生成client_crc
    data = k.read()
    print(data)

    print(k.server_crc, k.client_crc)
    # 支持crc64功能之前的文件，server_crc为None
    if k.server_crc is None:
        print("server crc is None, maybe server not support crc64")
    else:
      assert k.server_crc == k.client_crc, f"k.server_crc: {k.server_crc}, k.client_crc: {k.client_crc}"

# 下载为字符串后，计算本地crc64值，与服务端crc64值进行对比
def test_downloadObjectToString_checkingCrc64():
    k = bucket.get_key(key_name)
    s = k.get_contents_as_string().decode()
    client_crc = utils.compute_data_crc64(s)
    # client_crc = k.client_crc
    # 支持crc64功能之前的文件，server_crc为None
    if k.server_crc is None:
      print("server crc is None, maybe server not support crc64")
    else:
      print("client_crc:", client_crc, "server_crc:", k.server_crc)
      assert client_crc == k.server_crc, f"client_crc: {client_crc}, server_crc: {k.server_crc}"


# 下载到文件，计算本地文件crc64值，与服务端crc64值进行对比
def test_downloadObjectToFile_checkingCrc64():
    k = bucket.get_key(key_name)
    k.get_contents_to_filename('downloaded-' + key_name)
    # client_crc = utils.compute_file_crc64('downloaded-' + key_name)
    client_crc = k.client_crc
    # 支持crc64功能之前的文件，server_crc为None
    if k.server_crc is None:
      print("server crc is None, maybe server not support crc64")
    else:
      print("client_crc:", client_crc, "server_crc:", k.server_crc)
      assert client_crc == k.server_crc, f"client_crc: {client_crc}, server_crc: {k.server_crc}"


# 获取服务端CRC64值，与本地文件计算的CRC64值进行对比
def test_headObject_checkingCrc64():
    k = bucket.new_key(key_name)
    ret = k.set_contents_from_filename(file_path)
    assert ret.status == 200, f"ret.status: {ret.status} (expected 200)"

    k = bucket.get_key(key_name, validate=True)
    client_crc = utils.compute_file_crc64(file_path)
    assert client_crc == k.server_crc, f"client_crc: {client_crc}, k.server_crc: {k.server_crc}"


# copy文件后，获取目标文件的CRC64值，与本地文件计算的CRC64值进行对比
def test_copyObject_checkingCrc64():
    k = bucket.new_key(key_name)
    ret = k.set_contents_from_filename(file_path)
    assert ret.status == 200, f"ret.status: {ret.status} (expected 200)"

    target_key = bucket.copy_key("copy_test2", bucket_name, key_name)
    client_crc = utils.compute_file_crc64(file_path)
    assert client_crc == target_key.server_crc, f"client_crc: {client_crc}, target_key.server_crc: {target_key.server_crc}"


# 这么做，没什么意义。
# copy parts方法没有自动做crc校验（服务端保证），如果确实需要，可以手动计算本地源文件每个part的crc，
# 并通过合并得到最终的crc值，与服务端返回的crc进行对比
def test_multi_copy_checkingCrc64():
    k = bucket.new_key("test_multi")
    ret = k.set_contents_from_filename(file_path)
    assert ret.status == 200, f"ret.status: {ret.status} (expected 200)"

    copy_key = 'test_multi-copy'
    source_size = os.stat(file_path).st_size

    mp = bucket.initiate_multipart_upload(copy_key)

    chunk_size = 100 * 1024
    count = int(math.ceil(source_size * 1.0 / chunk_size * 1.0))

    help_list = []

    for i in range(count):
        start = chunk_size * i
        end = min(start + chunk_size - 1, source_size - 1)
        ret = mp.copy_part_from_key(bucket_name, "test_multi", i + 1, start, end)
        client_crc = utils.compute_file_crc64(file_path, start, end)
        part_info = ks3.multipart.Part()
        part_info.size = end - start + 1
        part_info.part_crc = client_crc
        help_list.append(part_info)
        assert client_crc == ret.server_crc, f"client_crc: {client_crc}, ret.server_crc: {ret.server_crc}"
    ret = mp.complete_upload()
    parts_crc = 0
    crc_obj = utils.Crc64(0)
    for p in help_list:
        parts_crc = crc_obj.combine(parts_crc, int(p.part_crc), p.size)
    assert str(parts_crc) == ret.ChecksumCRC64ECMA, f"parts_crc: {parts_crc}, ret.ChecksumCRC64ECMA: {ret.ChecksumCRC64ECMA}"


########## 手动计算crc，通过请求头校验的方式 ##########
# ks3服务端也提供了crc64校验的请求头，如果不开启自动校验crc的功能，可以通过传入请求头的方式让服务端进行crc的比对
def test_uploadObjectFromFile_checkingCrc64ByHeader():
    k = bucket.new_key(key_name)
    conn.enable_crc = False
    local_crc = utils.compute_file_crc64(file_path)
    crc_header = conn.provider.checksum_crc64ecma_header
    k.set_contents_from_filename(file_path, headers={crc_header: local_crc})


def test_uploadObjectFromString_checkingCrc64ByHeader():
    k = bucket.new_key(key_name)
    myContent = 'some string'
    conn.enable_crc = False
    local_crc = utils.compute_data_crc64(myContent)
    crc_header = conn.provider.checksum_crc64ecma_header
    k.set_contents_from_string(myContent, headers={crc_header: local_crc})


def test_multi_upload_checkingCrc64ByHeader():
    from filechunkio import FileChunkIO
    # 服务端不比对文件整体crc，可以在本地比对
    conn.enable_crc = False
    crc_header = conn.provider.checksum_crc64ecma_header
    crc_obj = utils.Crc64(0)
    final_crc = 0

    source_size = os.stat(file_path).st_size

    mp = bucket.initiate_multipart_upload("test_multi")

    chunk_size = 100 * 1024
    count = int(math.ceil(source_size * 1.0 / chunk_size * 1.0))

    for i in range(count):
        offset = chunk_size * i
        read_size = min(chunk_size, source_size - offset)
        with FileChunkIO(file_path, 'r', offset=offset, bytes=read_size) as fp:
            # 计算单个分片的crc
            local_part_crc = utils.compute_file_crc64(file_path, offset, offset + read_size - 1)
            # 将每一个分片的crc合并，最终将得到文件整体的crc值
            final_crc = crc_obj.combine(final_crc, int(local_part_crc), read_size)
            # 逐个上传分片
            mp.upload_part_from_file(fp, part_num=i + 1, headers={crc_header: local_part_crc})
    ret = mp.complete_upload()

    assert ret.ChecksumCRC64ECMA == str(final_crc), f"ret.ChecksumCRC64ECMA: {ret.ChecksumCRC64ECMA}, final_crc: {final_crc}"
    assert ret and ret.status == 200, f"ret: {ret}, ret.status: {getattr(ret, 'status', None)} (expected 200)"


if __name__ == '__main__':
    # test_downloadObject_checkingCrc64()
    # test_downloadObjectToString_checkingCrc64()
    # test_downloadObjectToFile_checkingCrc64()
    # test_headObject_checkingCrc64()
    # # test_copyObject_checkingCrc64()
    # # test_multi_copy_checkingCrc64()
    # test_uploadObjectFromFile_checkingCrc64ByHeader()
    # test_uploadObjectFromString_checkingCrc64ByHeader()
    # test_multi_upload_checkingCrc64ByHeader()
    test_downloadObjectToString_checkingCrc64()