import gradio as gr
import os,re
import numpy as np
import requests
import cv2
import base64
import json
from io import BytesIO
from PIL import Image

api_service_url = 'http://127.0.0.1:1234'
# api_service_url = 'http://i-2.gpushare.com:53412'

class ChangeFaceGui():
    def __init__(self):
        self.template_list = self.get_template_list()
        self.user_dict = self.get_user_dict()
        # self.tmp_dir = './tmp'
        # if not os.path.exists(self.tmp_dir): os.makedirs(self.tmp_dir)
        self.user_img_list = []
        # 获取户图片
        for item in self.user_dict:
            response = requests.get(item['face_img_url'])
            image_data = BytesIO(response.content)
            user_face_img = cv2.imdecode(np.frombuffer(image_data.read(), np.uint8), cv2.IMREAD_COLOR)
            self.user_img_list.append(user_face_img)
            item['face_img'] = user_face_img

            # img_path = f'{self.tmp_dir}/{item["user_id"]}.jpg'
            # if not os.path.exists(img_path):
            #     cv2.imwrite(f'{self.tmp_dir}/{item["user_id"]}.jpg', user_face_img)


    #请求得到模板图片的url
    def get_template_list(self):
        url = f"{api_service_url}/template/list"
        payload = {}
        headers = {}
        response = requests.request("POST", url, headers=headers, data=payload)
        print('请求模板照片成功！！')
        return(response.json()['data'])


    #请求得到用户图片及ID
    def get_user_dict(self):
        url = f"{api_service_url}/user/list"
        payload = {}
        headers = {}
        response = requests.request("POST", url, headers=headers, data=payload)
        user_dict = response.json()['data']
        print('请求用户照片成功！！')
        return(user_dict)


    #获取用户ID
    def get_user_id(self, user_img):
        user_small_img = cv2.resize(user_img, (100, 100))
        mse = []
        for item in self.user_dict:
            user_face_img = item['face_img']
            # cv2.imwrite(f'{self.tmp_dir}/{item["user_id"]}_list.jpg', user_face_img)
            user_face_img = cv2.resize(user_face_img, (100, 100))
            # 计算均方误差
            mse.append(np.mean((user_small_img - user_face_img) ** 2))

        # 找到均方差最小值对应的ID
        user_id = mse.index(min(mse))
        user_id = self.user_dict[user_id]['user_id']
        print('用户ID:', user_id)
        return user_id


    #虚拟试穿模块，输入是模特图和衣服图，输出是虚拟试穿的结果
    def take_photo(self, template_img, user_img):
        # if not os.path.exists(self.tmp_dir):os.makedirs(self.tmp_dir)

        if template_img is None or user_img is None:
            return None

        # 使用 Pillow 加载图像
        template_img = Image.open(template_img)
        user_img = Image.open(user_img)

        # 将 Pillow 图像转换为 OpenCV 格式（BGR）
        template_img = cv2.cvtColor(np.array(template_img), cv2.COLOR_RGB2BGR)
        user_img = cv2.cvtColor(np.array(user_img), cv2.COLOR_RGB2BGR)

        # 将模板图片转换为base64格式
        retval, template_bytes = cv2.imencode('.jpg', template_img)
        encoded_image = base64.b64encode(template_bytes).decode('utf-8')
        user_id = self.get_user_id(user_img)
        url = f"{api_service_url}/user/generate"
        payload = json.dumps({
            "user_id": user_id,
            "base_img": encoded_image
        })
        headers = {
            'Content-Type': 'application/json'
        }
        print('请求api_service.py发送请求！！')
        response = requests.request("POST", url, headers=headers, data=payload)
        print('请求api_service.py发送请求成功！！')
        if response.status_code != 200:
            raise RuntimeError(f"Failed to send request to API service. Status code: {response.status_code}")

        ret_image_b64 = response.json().get('generate_photo_b64')

        if ret_image_b64 is None:
            raise RuntimeError(f"ret image failed!")

        image_array = np.frombuffer(base64.b64decode(ret_image_b64), np.uint8)
        result_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
        result_image = cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB)
        return result_image


    def start_gui(self):
        user_img_list = []
        for item in self.user_img_list:
            img = cv2.cvtColor(item, cv2.COLOR_BGR2RGB)
            user_img_list.append(img)

        with gr.Blocks() as demo:
            with gr.Row():
                gr.Markdown("# 数字力场效果展示")

            # 换脸
            with gr.Tab("数字写真"):
                with gr.Row():
                    gr.Markdown("# 数字写真")
                with gr.Row():
                    with gr.Column():
                        #选择模板
                        template_img = gr.Image(label="模版", sources='upload', min_width=384, width=384, height=384, type="filepath", value=self.template_list[0],interactive=True)
                        example_template = gr.Examples(
                            inputs=template_img,
                            examples_per_page=12,
                            examples= self.template_list)
                    with gr.Column():
                        # 选择用户
                        user_img = gr.Image(label="用户", sources='upload', min_width=384, width=384, height=384, type="filepath", value= user_img_list[0],interactive=False)
                        example_user = gr.Examples(
                            inputs=user_img,
                            examples_per_page=12,
                            examples=user_img_list)

                    with gr.Column():
                        output_img = gr.Image(label="结果展示", type="numpy", height=576, width=384)
                with gr.Column():
                    run_button = gr.Button(value="提交")
                    run_button.click(self.take_photo, inputs=[template_img, user_img], outputs=[output_img])

            #             #换衣服
            # with gr.Tab("换衣"):
            #     with gr.Row():
            #         gr.Markdown("# 换衣")
            #     text_button = gr.Button("提交")
            # # 换发型
            # with gr.Tab("换发型"):
            #     with gr.Row():
            #         gr.Markdown("# 换发型")
            #     text_button = gr.Button("提交")

        demo.launch(server_name='0.0.0.0', server_port=8080)


if __name__ == '__main__':
    demo = ChangeFaceGui()
    demo.start_gui()