import insightface
import os
import onnxruntime
import cv2
import gfpgan
import tempfile
import time
import gradio as gr
import sys
from torchvision.transforms import functional
from PIL import Image

# 참조코드에서 사용된 모듈 임포트 수정
sys.modules["torchvision.transforms.functional_tensor"] = functional

class Predictor:
    def __init__(self):
        self.setup()

    def setup(self):
        os.makedirs('models', exist_ok=True)
        os.chdir('models')
        if not os.path.exists('GFPGANv1.4.pth'):
            os.system(
                'wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth'
            )
        if not os.path.exists('inswapper_128.onnx'):
            os.system(
                'wget https://huggingface.co/ashleykleynhans/inswapper/resolve/main/inswapper_128.onnx'
            )
        os.chdir('..')  # 디렉토리 변경 완료

        """💎 Load the model into memory to make running multiple predictions efficient"""
        self.face_swapper = insightface.model_zoo.get_model('models/inswapper_128.onnx',
                                                            providers=onnxruntime.get_available_providers())
        # self.face_swapper.prepare(ctx_id=0, det_size=(640, 640))  # 이 줄을 제거합니다.

        self.face_enhancer = gfpgan.GFPGANer(model_path='models/GFPGANv1.4.pth', upscale=1)
        self.face_analyser = insightface.app.FaceAnalysis(name='buffalo_l')
        self.face_analyser.prepare(ctx_id=0, det_size=(640, 640))

    def get_face_image(self, img_data, face):
        # 얼굴 영역을 잘라내기
        x1, y1, x2, y2 = [int(coord) for coord in face.bbox]
        face_img = img_data[y1:y2, x1:x2]
        return face_img

    def predict(self, input_image_path, swap_image_path):
        """🧶 Run a single prediction on the model"""
        try:
            frame = cv2.imread(input_image_path)
            if frame is None:
                print("❌ Target image could not be read.")
                return None
            analysed = self.face_analyser.get(frame)
            if not analysed:
                print("❌ No face found in target image.")
                return None
            face = max(analysed, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]))
            target_face_img = self.get_face_image(frame, face)

            swap_frame = cv2.imread(swap_image_path)
            if swap_frame is None:
                print("❌ Swap image could not be read.")
                return None
            swap_analysed = self.face_analyser.get(swap_frame)
            if not swap_analysed:
                print("❌ No face found in swap image.")
                return None
            swap_face = max(swap_analysed, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]))
            swap_face_img = self.get_face_image(swap_frame, swap_face)

            # 얼굴 교체 수행
            result = self.face_swapper.get(frame, face, swap_face, paste_back=True)

            # 얼굴 향상 수행
            _, _, result = self.face_enhancer.enhance(
                result,
                paste_back=True
            )
            out_path = os.path.join(tempfile.mkdtemp(), f"{str(int(time.time()))}.jpg")
            cv2.imwrite(out_path, result)
            return out_path
        except Exception as e:
            print(f"{e}")
            return None

# Predictor 클래스 인스턴스 생성
predictor = Predictor()

# CSS 및 테마 설정
css = """
/* "Swap Faces" 버튼 스타일 */
button#swap-button {
    background-color: #FB923C !important; /* 주황색 배경 */
    color: white !important; /* 흰색 글씨 */
}
/* "이미지 다운로드 (JPG)" 버튼 스타일 */
button#download-button {
    background-color: #FB923C !important; /* 주황색 배경 */
    color: white !important; /* 흰색 글씨 */
}
/* 필요에 따라 추가적인 스타일을 여기에 작성할 수 있습니다 */
"""

demo_theme = gr.themes.Soft(
    primary_hue=gr.themes.Color(
        c50="#FFF7ED",
        c100="#FFEDD5",
        c200="#FED7AA",
        c300="#FDBA74",
        c400="#FB923C",
        c500="#F97316",
        c600="#EA580C",
        c700="#C2410C",
        c800="#9A3412",
        c900="#7C2D12",
        c950="#431407",
    ),
    secondary_hue="zinc",
    neutral_hue="zinc",
    font=("Pretendard", "sans-serif")
)

# JPG 다운로드 기능 구현
def save_as_jpg(file_path):
    try:
        if file_path is None:
            return None
        # 파일 경로를 받아 PIL 이미지로 변환
        img = Image.open(file_path)
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
            img.save(tmp, format="JPEG")
            tmp_path = tmp.name
        return tmp_path  # 파일 경로 반환
    except Exception as error:
        print(f"Error saving as JPG: {error}")
        return None

# Clear 함수: 입력 및 출력 초기화
def clear_all():
    return [None, None, None]

# Gradio Interface 구성
with gr.Blocks(theme=demo_theme, css=css) as demo:
    with gr.Row():
        # 왼쪽 섹션: 입력
        with gr.Column(scale=1):
            target_image = gr.Image(
                type="filepath",
                label="얼굴을 변경할 이미지"
            )
            swap_image = gr.Image(
                type="filepath",
                label="교체할 얼굴"
            )
            swap_button = gr.Button("얼굴 교체", elem_id="swap-button")
            clear_button = gr.Button("리셋하기")
        
        # 오른쪽 섹션: 출력
        with gr.Column(scale=1):
            result_image = gr.Image(
                type="filepath",
                label="결과 이미지"
            )
            download_jpg_button = gr.Button("JPG로 변환하기", elem_id="download-button")
            download_file = gr.File(label="JPG 이미지 다운받기")
    
    # 버튼 클릭 시 예측 함수 호출
    swap_button.click(
        fn=predictor.predict,
        inputs=[target_image, swap_image],
        outputs=result_image
    )

    # 리셋하기 버튼 클릭 시 입력 및 출력 이미지 초기화
    clear_button.click(
        fn=clear_all,
        inputs=None,
        outputs=[target_image, swap_image, result_image]
    )

    # JPG 다운로드 버튼 클릭 시 파일 생성
    download_jpg_button.click(
        fn=save_as_jpg,
        inputs=result_image,
        outputs=download_file
    )

# Gradio Interface 실행
demo.launch()