import oss2, os, sys, shutil, cv2, numpy
from PIL import Image
from datetime import datetime
import config
from threading import Thread

def downfile(key, localfile):
    oss_client.get_object_to_file(key, localfile)
    try:
        im = Image.open(localfile).load()
    except Exception as e:
        print(e)
        print('删除错误结构的图片文件:', localfile)
        os.remove(localfile)

def down_photos_from_oss(path, pid):
    if os.path.exists(path): return
    os.system('mkdir ' + path + '\photo1 '  + path + '\photo2')
    # 根据前缀获取文件列表
    prefix = f'photos/{pid}/'
    filelist = oss2.ObjectIteratorV2(oss_client, prefix=prefix)
    for file in filelist:
        filename = file.key.split('/')[-1]
        if filename in ('16_1.jpg', '26_1.jpg', '36_1.jpg', '46_1.jpg', '56_1.jpg', '66_1.jpg', '76_1.jpg', '86_1.jpg', '16_8.jpg', '26_8.jpg', '36_8.jpg', '46_8.jpg', '56_8.jpg', '66_8.jpg', '76_8.jpg', '86_8.jpg'): continue
        if filename.endswith('_1.jpg'):
            localfile = os.path.join(path, 'photo1', filename)
        else:
            localfile = os.path.join(path, 'photo2', filename)

        print('正在下载:', localfile)
        t = Thread(target=downfile, args=(file.key, localfile))
        t.start()

def down_xmps_from_oss(psid):
    dest_path = os.path.join(config_path, 'xmps', psid)
    if os.path.exists(dest_path): return
    
    os.makedirs(dest_path)
    os.makedirs(os.path.join(dest_path, 'mesh'))
    os.makedirs(os.path.join(dest_path, 'texture'))

    prefix = f'xmps/{psid}/mesh/'
    filelist = oss2.ObjectIteratorV2(oss_client, prefix=prefix)
    for file in filelist:
        if file.key.endswith('.xmp'): 
            filename = file.key.split('/')[-1]
            print('正在下载:', file.key)
            oss_client.get_object_to_file(file.key, os.path.join(dest_path, 'mesh', filename))

    prefix = f'xmps/{psid}/texture/'
    filelist = oss2.ObjectIteratorV2(oss_client, prefix=prefix)
    for file in filelist:
        if file.key.endswith('.xmp'):
            filename = file.key.split('/')[-1]
            print('正在下载:', file.key)
            oss_client.get_object_to_file(file.key, os.path.join(dest_path, 'texture', filename))

    print('下载完成')

def adjust_lab(img, a=7, b=5):
    img = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
    img[:, :, 1] = img[:, :, 1] - a
    img[:, :, 2] = img[:, :, 2] + b
    img = cv2.cvtColor(img, cv2.COLOR_LAB2BGR)
    return img

def enhance_photos(pid):
    path = os.path.join(workdir, pid, 'photo1')
    for filename in os.listdir(path):
        if filename.endswith('.jpg'):
            print('正在处理:', filename)
            image = cv2.imread(os.path.join(path, filename))
            image = adjust_lab(image)
            image = adjust_gamma(image, 1, 2.2)
            cv2.imwrite(os.path.join(path, filename), image)

    path = os.path.join(workdir, pid, 'photo2')
    for filename in os.listdir(path):
        if filename.endswith('.jpg'):
            print('正在处理:', filename)
            image = cv2.imread(os.path.join(path, filename))
            image = adjust_lab(image)
            cv2.imwrite(os.path.join(path, filename), image)

def adjust_gamma(img, c=1, g=2.2):
    img = img.astype(float)
    out = img.copy()
    out /= 255.
    out = (1/c * out) ** (1/g)
    out *= 255
    out = out.astype(numpy.uint8)
    return out

def find_color_area(img, lower_color, upper_color):
    color = img.copy()
    color = cv2.cvtColor(color, cv2.COLOR_BGR2GRAY)
    color = color[2800:, :]
    mask = cv2.inRange(color, lower_color, upper_color)
    area = cv2.countNonZero(mask)
    return round(area / (color.shape[0] * color.shape[1]), 4)

def clean_photo2(psid, pid):
    path = os.path.join(workdir, pid, 'photo2')
    print('正在处理psid:', psid, '的颜色配置：', config.ps_lights.get(psid))
    if config.ps_lights.get(psid) is None:
        ps_lights = config.ps_lights['0']
    else:
        ps_lights = config.ps_lights[psid]

    for filename in os.listdir(path):
        if not filename.endswith('_8.jpg'): continue
        img = cv2.imread(os.path.join(path, filename))
        white_area = find_color_area(img.copy(), ps_lights['lower_white'], ps_lights['upper_white'])
        gray_area = find_color_area(img.copy(), ps_lights['lower_gray'], ps_lights['upper_gray'])
        print('file:', os.path.join(path, filename), 'white_area:', white_area, 'gray_area:', gray_area)
        if filename == '25_8.jpg':
            print('file:', os.path.join(path, filename), 'white_area:', white_area, 'gray_area:', gray_area)
        if white_area < gray_area or gray_area > ps_lights['gray_threshold']:
            print('删除亮度异常图片:', os.path.join(path, filename), '亮度:', white_area, '灰度:', gray_area)
            os.remove(os.path.join(path, filename))

def diff_time(start_time, end_time):
    d = end_time - start_time
    m = d.microseconds
    return str(d.seconds // 60) + ":" + str(d.seconds % 60) + ":" + str(m)[:-3]

def convert_to_gray(pid):
    path = os.path.join(workdir, pid, 'photo1')
    for filename in os.listdir(path):
        if filename.endswith('.jpg'):
            print('正在处理:', filename)
            image = Image.open(os.path.join(path, filename))
            image = image.convert('L')
            image = image.rotate(180)
            image.save(os.path.join(path, filename))
    
    path = os.path.join(workdir, pid, 'photo2')
    for filename in os.listdir(path):
        if filename.endswith('.jpg'):
            print('正在处理:', filename)
            image = Image.open(os.path.join(path, filename))
            image = image.rotate(180)
            image.save(os.path.join(path, filename))

if __name__ == '__main__':
    workdir = 'z:\\'
    AccessKeyId = 'LTAI5tSReWm8hz7dSYxxth8f'
    AccessKeySecret = '8ywTDF9upPAtvgXtLKALY2iMYHIxdS'
    Endpoint = 'oss-cn-shanghai.aliyuncs.com'
    Bucket = 'suwa3d-securedata'
    oss_client = oss2.Bucket(oss2.Auth(AccessKeyId, AccessKeySecret), Endpoint, Bucket)
    config_path = 'D:\\apps\\config\\'

    if len(sys.argv) == 1:
        print('usage: python make_3d.py [pid]')
        exit(1)
    else:
        pid = sys.argv[1]
        path = os.path.join(workdir, pid)

    start_time = datetime.now()
    down_photos_from_oss(path, pid)

    start_clean_time = datetime.now()
    print('下载花费时间:', diff_time(start_time, start_clean_time))
    # clean_photo2(0, pid)
    
    start_enhance_time = datetime.now()
    print('清理无效图片花费时间:', diff_time(start_clean_time, start_enhance_time))
    # enhance_photos(pid)
    # convert_to_gray(pid)

    end_time = datetime.now()
    print('优化图片花费时间:', diff_time(start_enhance_time, end_time))
    print('总花费时间:', diff_time(start_time, end_time))
    # make 3d