#!/usr/bin/python
import oss2, os, sys, redis, time, cv2, numpy, piexif, requests, shlex, subprocess, datetime, shutil, json
from threading import Thread
from matplotlib import pyplot as plt
from PIL import Image, ImageEnhance
import config

def find_blender_bin_path():
    base_path = 'C:\\Program Files\\Blender Foundation\\'
    if os.path.exists(base_path):
        for dir in os.listdir(base_path):
            if dir.startswith('Blender'):
                blender_bin_path = base_path + dir + '\\blender.exe'
                return f'"{blender_bin_path}"'
    else:
        print('未找到blender安装目录')
        exit(1)

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 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_from_oss(pid):
    path = os.path.join(workdir, pid)
    if os.path.exists(path): return
    os.makedirs(path)
    os.mkdir(os.path.join(path, 'photo1'))
    os.mkdir(os.path.join(path, 'photo2'))
    # os.mkdir(os.path.join(path, 'photo2_bad'))
    
    psid = getPSid(pid)
    # 根据前缀获取文件列表
    prefix = f'photos/{pid}/'
    filelist = oss2.ObjectIteratorV2(oss_client, prefix=prefix)
    for file in filelist:
        filename = file.key.split('/')[-1]
        if psid == '1' and filename in ('113_8.jpg', '113_1.jpg'): continue
        if psid == '9' and filename in ('135_8.jpg', '135_1.jpg'): continue
        if psid == '16' and filename in ('92_8.jpg', '92_1.jpg'): continue
        if psid == '25' and filename in ('13_8.jpg', '13_1.jpg'): continue
        if psid == '34' and filename in ('154_8.jpg', '154_1.jpg'): continue
        if psid == '44' and filename in ('31_8.jpg', '31_1.jpg'): continue
        if psid == '63' and filename in ('123_8.jpg', '123_1.jpg'): continue
        if psid == '70' and filename in ('73_8.jpg', '73_1.jpg'): continue
        print('正在下载:', file.key)
        if filename.endswith('_1.jpg'):
            localfile = os.path.join(path, 'photo1', filename)
        else:
            localfile = os.path.join(path, 'photo2', filename)
        oss_client.get_object_to_file(file.key, localfile)
     
    print('下载完成')

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_brightness(image, brightness_factor):
    if brightness_factor == 1 or brightness_factor == 0 :
        return image
    enhancer = ImageEnhance.Brightness(image)
    adjusted_image = enhancer.enhance(brightness_factor)
    return adjusted_image

def adjust_saturation(image, saturation_factor):
    if saturation_factor == 1:
        return image
    enhancer = ImageEnhance.Color(image)
    adjusted_image = enhancer.enhance(saturation_factor)
    return adjusted_image

def adjust_temperature(image, temperature_factor):
    if temperature_factor == 1:
        return image
    r, g, b = image.split()
    r = r.point(lambda i: i * temperature_factor)
    adjusted_image = Image.merge("RGB", (r, g, b))
    return adjusted_image

def rotate_image(image):
    # 检查图像的EXIF数据是否包含方向信息
    try:
        exif = image._getexif()
        orientation = exif.get(0x0112)
    except:
        orientation = None

    # 根据方向信息旋转图像
    if orientation == 3:
        image = image.rotate(180, expand=True)
    elif orientation == 6:
        image = image.rotate(270, expand=True)
    elif orientation == 8:
        image = image.rotate(90, expand=True)
    return image

def get_ps_adjust_photo_para(psid):
    res = requests.get(get_ps_adjust_photo_para_url, params={'id': psid})
    print(res.json())
    paras = res.json()['data']
    brightness_factor, saturation_factor, temperature_factor = float(paras['brightness']), float(paras['saturation']), float(paras['colorTemperature'])
    return brightness_factor, saturation_factor, temperature_factor

def adjust_photos(pid):
    if not os.path.exists(os.path.join(workdir, pid, 'photo2')):
        print(f"Directory {os.path.join(workdir, pid, 'photo2')} does not exist")
        return False
    
    psid = getPSid(pid)
    brightness_factor, saturation_factor, temperature_factor = get_ps_adjust_photo_para(psid)

    if (brightness_factor == 1 and saturation_factor == 1 and temperature_factor == 1):
        print("No need to adjust")
        return False
    os.makedirs(os.path.join(workdir, pid, 'photo3'), exist_ok=True)

    for filename in os.listdir(os.path.join(workdir, pid, 'photo2')):
        if filename.endswith(".jpg"):
            print(f"Adjusting {filename}:brightness={brightness_factor}, saturation={saturation_factor}, temperature={temperature_factor}")
            image = Image.open(os.path.join(workdir, pid, 'photo2', filename))
            image = rotate_image(image)

            brightened_image = adjust_brightness(image, brightness_factor)
            saturated_image = adjust_saturation(brightened_image, saturation_factor)
            adjusted_image = adjust_temperature(saturated_image, temperature_factor)

            adjusted_image.save(os.path.join(workdir, pid, 'photo3', filename))
    print("Done")
    return True

def getPSid(pid):
    res = requests.get(get_psid_url, params={'pid': pid})
    print('get_psid_url:', res.url)
    print('res:', res.text)
    res = json.loads(res.text)
    return str(res['data'])

def getHeadCount(pid):
    res = requests.get(get_printinfo_url, params={'id': pid})
    print('get_printinfo_url:', res.url)
    print('res:', res.text)
    if res.status_code != 200:
        print('获取人数失败，程序退出')
        exit(1)
    res = json.loads(res.text)
    return res['data']['headcount']

def aliyun_face(pid, camera_id):
    high = False
    style = 'imm/detectface'
    objectkey = f'photos/{pid}/photo2/{camera_id}_8.jpg'
    try:
        res = oss_client.get_object(objectkey, process=style)
    except oss2.exceptions.NoSuchKey:
        print('没有找到文件:', objectkey)
        return high
    res = json.loads(res.read())
    if res['success']:
        if res['Faces'] is None:
            print('no face')
            return None
        else:
            print('faces num:', len(res['Faces']))
        for face in res['Faces']:
            print('-' * 20)
            print('face_id:', face['FaceId'])
            print('gender:', face['Gender'])
            print('age:', face['Age'])

            if face['Gender'] == 'FEMALE' and face['Age'] < 22: high = True
            if face['Gender'] == 'MALE' and face['Age'] < 15: high = True
    else:
        print('face detect failed...')
    return high

def set_photos_join_type(pid, photoN, mesh = '0', texture='0'):
    photoN_path = os.path.join(workdir, pid, photoN)
    for xmp in os.listdir(photoN_path):
        if xmp.endswith('.xmp'):
            xmp_path = os.path.join(photoN_path, xmp)
            with open(xmp_path, 'r') as f:
                lines = f.readlines()
                lines = [line.replace('xcr:InMeshing="0"', f'xcr:InMeshing="{mesh}"') for line in lines]
                lines = [line.replace('xcr:InMeshing="1"', f'xcr:InMeshing="{mesh}"') for line in lines]
                lines = [line.replace('xcr:InTexturing="0"', f'xcr:InTexturing="{texture}"') for line in lines]
                lines = [line.replace('xcr:InTexturing="1"', f'xcr:InTexturing="{texture}"') for line in lines]
                with open(xmp_path, 'w') as f:
                    f.writelines(lines)

def rc_make(pid, psid):
    status = True
    # 建模
    start = datetime.datetime.now()

    xmps_mesh_path = config_path + f'xmps\\{psid}\\mesh\\'
    xmps_texture_path = config_path + f'xmps\\{psid}\\texture\\'

    input_path = os.path.join(workdir, pid)
    output_path = os.path.join(input_path, 'output')
    ImagesGeometry = os.path.join(input_path, "photo1")
    ImagesTexture = os.path.join(input_path, "photo2")
    photo3 = os.path.join(input_path, "photo3")

    ps_region = psid + '.rcbox'

    shutil.rmtree(output_path, ignore_errors=True)
    os.system('mkdir ' + output_path)

    # ps_fix_list = ('10', '12', '13', '16', '18', '19', '20', '22', '23', '29', '30', '31')
    os.system('xcopy /Y ' + xmps_mesh_path + '* ' + ImagesGeometry)
    os.system('xcopy /Y ' + xmps_texture_path + '* ' + ImagesTexture)

    simplify_value = str(getHeadCount(pid) * 100 * 10000)
    
    if adjust_photos(pid): # 已调整照片，产生出photo3图片组，让photo1和photo2参与建模，photo3仅参与贴图
        # if not os.path.exists(os.path.join(photo3, '11_8.xmp')):
        #     os.system('xcopy /Y ' + xmps_texture_path + '* ' + photo3)
        set_photos_join_type(pid, 'photo1', mesh='1', texture='0')
        set_photos_join_type(pid, 'photo2', mesh='1', texture='0')
        set_photos_join_type(pid, 'photo3', mesh='0', texture='1')
        cmd = rcbin + \
            ' -disableOnlineCommunication ' + \
            ' -silent ' + output_path + \
            ' -addFolder "' + ImagesGeometry + '"' + \
            ' -set "sfmMaxFeaturesPerMpx=20000" -set "sfmMaxFeaturesPerImage=200000" -set "sfmImagesOverlap=High" -set "sfmMaxFeatureReprojectionError=2" ' + \
            ' -align ' + \
            ' -addFolder "' + ImagesTexture + '"' + \
            ' -align ' + \
            ' -addFolder "' + photo3 + '"' + \
            ' -align -align' + \
            ' -setReconstructionRegion "' + config_path + ps_region + '" ' + \
            ' -mvs ' + \
            ' -modelSelectMaximalConnectedComponent -modelInvertSelection -modelRemoveSelectedTriangles -clean ' + \
            ' -simplify ' + simplify_value + ' ' + \
            ' -smooth -unwrap ' + \
            ' -calculateTexture ' + \
            ' -renameModel ' + pid + ' ' + \
            ' -exportModel ' + pid + ' "' + output_path + '\\' + pid + '.obj" "' + config_path + 'ModelExportParams102.xml" ' + \
            ' -quit'
    else: # 未调整照片，仍然使用photo1和photo2建模，photo2参与贴图模式
        set_photos_join_type(pid, 'photo1', mesh='1', texture='0')
        set_photos_join_type(pid, 'photo2', mesh='1', texture='1')
        cmd = rcbin + \
            ' -disableOnlineCommunication ' + \
            ' -silent ' + output_path + \
            ' -addFolder "' + ImagesGeometry + '"' + \
            ' -set "sfmMaxFeaturesPerMpx=20000" -set "sfmMaxFeaturesPerImage=200000" -set "sfmImagesOverlap=High" -set "sfmMaxFeatureReprojectionError=2" ' + \
            ' -align ' + \
            ' -addFolder "' + ImagesTexture + '"' + \
            ' -align -align' + \
            ' -setReconstructionRegion "' + config_path + ps_region + '" ' + \
            ' -mvs ' + \
            ' -modelSelectMaximalConnectedComponent -modelInvertSelection -modelRemoveSelectedTriangles -clean ' + \
            ' -simplify ' + simplify_value + ' ' + \
            ' -smooth -unwrap ' + \
            ' -calculateTexture ' + \
            ' -renameModel ' + pid + ' ' + \
            ' -exportModel ' + pid + ' "' + output_path + '\\' + pid + '.obj" "' + config_path + 'ModelExportParams102.xml" ' + \
            ' -quit'

    print(cmd)
    cmd = shlex.split(cmd)
    res = subprocess.run(cmd)
    print(res)
    
    end = datetime.datetime.now()
    print('Time: ' + diff_time(start, end))

    return status

def fix_mtl_file(pid):
    path = os.path.join(workdir, pid)
    mtl_file = os.path.join(path, 'output', pid + '.mtl')
    os.rename(os.path.join(path, 'output', pid + '_u1_v1.jpg'), os.path.join(path, 'output', pid + '.jpg'))
    with open(os.path.join(path, mtl_file), 'r') as f:
        lines = f.readlines()
        lines = [line.replace('_u1_v1.jpg', '.jpg') for line in lines]
        with open(os.path.join(path, mtl_file), 'w') as f:
            f.writelines(lines)

def delete_lines_in_file(filename, count):
    with open(filename, 'r') as f:
        lines = f.readlines()
        lines = lines[count:]
        with open(filename, 'w') as f:
            f.writelines(lines)

def find_gray(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gray = gray[2800:, :]
    hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
    plt.plot(hist)

def find_color_area(img, lower_color, upper_color):
    color = cv2.cvtColor(img, 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 auto_make_3d(r):
    while True:
        try:
            if r.llen('model:make10') == 0:
                time.sleep(5)
                continue
        except Exception as e:
            print(e)
            time.sleep(5)
            r = redis.Redis(host='106.14.158.208', password='kcV2000', port=6379, db=6)
            continue 
        pid = r.lpop('model:make10')
        if pid is None: continue
        pid = pid.decode('utf-8')
        psid = getPSid(pid)
        down_xmps_from_oss(psid)

        print(time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time())), f'pid:{pid} 模型重建 start')
        res = requests.post(update_status_modeling_url, data={'id': pid})
        print('更新建模中状态：', res.text)

        # 从oss下载图片
        down_from_oss(pid)

        # 调用rc建模
        if rc_make(pid, psid):
            fix_mtl_file(pid)
            os.system(f'{blenderbin} -b -P d:\\apps\\blender\\autofix.py -- {pid}')
            os.system(f'gltfpack -c -i {os.path.join(workdir, pid, "output", pid + ".obj")} -o {os.path.join(workdir, pid, "output", pid + ".glb")}')
            print(time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time())), f'pid:{pid} 建模完成，开始上传到oss')
            print('建模成功: ' + pid, '开始上传模型文件...')
            oss_client = oss2.Bucket(oss2.Auth(AccessKeyId, AccessKeySecret), Endpoint, Bucket)
            oss_client.put_object_from_file(f'objs/auto/{pid}/{pid}.obj.rcInfo', os.path.join(workdir, pid, 'output', pid + '.obj.rcInfo'))
            oss_client.put_object_from_file(f'objs/auto/{pid}/{pid}.obj', os.path.join(workdir, pid, 'output', pid + '.obj'))
            oss_client.put_object_from_file(f'objs/auto/{pid}/{pid}.mtl', os.path.join(workdir, pid, 'output', pid + '.mtl'))
            oss_client.put_object_from_file(f'objs/auto/{pid}/{pid}.jpg', os.path.join(workdir, pid, 'output', pid + '.jpg'))
            oss_client.put_object_from_file(f'glbs/auto/{pid}.glb', os.path.join(workdir, pid, 'output', pid + '.glb'))
            res = requests.post(update_status_modelsuccess_url, data={'id': pid})
            print('上传完成更新建模成功状态：', res.text)
            print(time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()), f'pid:{pid} 建模完成,上传模型及更新状态 end')
        else:
            print(time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()), f'pid:{pid} 建模失败')
            # 建模失败，放回队列，微信通知技术人员介入处理
            r.lpush('model:make10', pid)
            res = requests.post(update_status_modelfailed_url, data={'id': pid})

        # 清空本地缓存
        shutil.rmtree(os.path.join(workdir, pid), ignore_errors=True)
        shutil.rmtree(temp_path, ignore_errors=True)

        time.sleep(1)

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

    r = redis.Redis(host='106.14.158.208', password='kcV2000', port=6379, db=6)
    update_status_modeling_url = 'https://mp.api.suwa3d.com/api/customerP3dLog/toModeling'
    update_status_modelsuccess_url = 'https://repair.api.suwa3d.com/api/modelRepairOrder/toModelMakeSucceed'
    update_status_modelfailed_url = 'https://mp.api.suwa3d.com/api/customerP3dLog/toModelMakeFailed'
    get_psid_url = 'https://mp.api.suwa3d.com/api/customerP3dLog/photoStudio'
    get_printinfo_url = 'https://mp.api.suwa3d.com/api/customerP3dLog/printInfo'
    update_status_printstatus_url = 'https://mp.api.suwa3d.com/api/customerP3dLog/updateBuildPrintModelStatus'
    get_ps_adjust_photo_para_url = 'https://mp.api.suwa3d.com/api/equipment/configForColor' # 获取ps调整图片参数的接口
    workdir = 'D:\\'
    rcbin = '"C:\\Program Files\\Capturing Reality\\RealityCapture\\RealityCapture.exe"'
    blenderbin = find_blender_bin_path()
    config_path = 'D:\\apps\\config10\\'
    temp_path = 'D:\\RealityCapture.Temp\\'

    print(time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()), '3D重建程序 start ', )

    auto_make_3d(r)