import torch
from typing import List
from insightface.app.common import Face

from gfpgan.utils import GFPGANer
from .typing import Frame # 你已有定义
from auto_media_publisher.config.conf_base import FACE_ENHANCER_MODEL_PATH

SIMILAR_FACE_DISTANCE: float = 0.45     # < 0.4 同一人 0.4 -0.6 大概是同一人 >0.6 不是同一个


class FaceEnhancer:
    def __init__(self):
        self.device = self._select_device()
        self._enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=self.device)

    def _select_device(self):
        return 'cpu'
        if torch.cuda.is_available():
            return 'cuda'
        elif torch.backends.mps.is_available():
            return 'mps'
        else:
            return 'cpu'
        
    def _enhance_face(self,target_face: Face, temp_frame: Frame) -> Frame:
        start_x, start_y, end_x, end_y = map(int, target_face['bbox'])
        padding_x = int((end_x - start_x) * 0.5)
        padding_y = int((end_y - start_y) * 0.5)
        start_x = max(0, start_x - padding_x)
        start_y = max(0, start_y - padding_y)
        end_x = max(0, end_x + padding_x)
        end_y = max(0, end_y + padding_y)
        temp_face = temp_frame[start_y:end_y, start_x:end_x]
        if temp_face.size:
            _, _, temp_face = self._enhancer.enhance(
                temp_face,
                paste_back=True
            )
            temp_frame[start_y:end_y, start_x:end_x] = temp_face
        return temp_frame


    def process_frame(self,temp_frame: Frame,many_faces:List) -> Frame:
        #many_faces = self._get_many_faces(temp_frame)  check face_analyser.py
        if many_faces:
            for target_face in many_faces:
                temp_frame = self._enhance_face(target_face, temp_frame)
        return temp_frame