import cv2
import mediapipe as mp
import os
from gradio_client import Client
# from test_image_fusion import Test
# from test_image_fusion import Test
from test_image import Test
import numpy as np



from PIL import Image
import numpy as np
import cv2

# client = Client("https://tbvl-real-and-fake-face-detection.hf.space/--replicas/40d41jxhhx/")

data = 'faceswap'
dct = 'fft'


# testet = Test(model_paths = [f"weights/{data}-hh-best_model.pth",
#                             f"weights/{data}-fft-best_model.pth"],
#                             multi_modal = ['hh', 'fft'])

testet = Test(model_path =f"weights/{data}-hh-best_model.pth",
                            multi_modal ='hh')

# Initialize MediaPipe Face Detection
mp_face_detection = mp.solutions.face_detection
mp_drawing = mp.solutions.drawing_utils
face_detection = mp_face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.35)

# Create a directory to save the cropped face images if it does not exist
save_dir = "cropped_faces"
os.makedirs(save_dir, exist_ok=True)

# def detect_and_label_faces(image_path):


# Function to crop faces from a video and save them as images
# def crop_faces_from_video(video_path):
#     # Read the video
#     cap = cv2.VideoCapture(video_path)
#     frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
#     frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
#     fps = int(cap.get(cv2.CAP_PROP_FPS))
#     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
#     # Define the codec and create VideoWriter object
#     out = cv2.VideoWriter(f'output_{real}_{data}_fusion.avi', cv2.VideoWriter_fourcc('M','J','P','G'), fps, (frame_width, frame_height))

#     if not cap.isOpened():
#         print("Error: Could not open video.")
#         return
# Convert PIL Image to NumPy array for OpenCV
def pil_to_opencv(pil_image):
    open_cv_image = np.array(pil_image)
    # Convert RGB to BGR for OpenCV
    open_cv_image = open_cv_image[:, :, ::-1].copy()
    return open_cv_image

# Convert OpenCV NumPy array to PIL Image
def opencv_to_pil(opencv_image):
    # Convert BGR to RGB
    pil_image = Image.fromarray(opencv_image[:, :, ::-1])
    return pil_image




def detect_and_label_faces(frame):
    frame = pil_to_opencv(frame)


    print(type(frame))
    # Convert the frame to RGB
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    # Perform face detection
    results = face_detection.process(frame_rgb)

    # If faces are detected, crop and save each face as an image
    if results.detections:
        for face_count,detection in enumerate(results.detections):
            bboxC = detection.location_data.relative_bounding_box
            ih, iw, _ = frame.shape
            x, y, w, h = int(bboxC.xmin * iw), int(bboxC.ymin * ih), int(bboxC.width * iw), int(bboxC.height * ih)
            # Crop the face region and make sure the bounding box is within the frame dimensions
            crop_img = frame[max(0, y):min(ih, y+h), max(0, x):min(iw, x+w)]
            if crop_img.size > 0:
                face_filename = os.path.join(save_dir, f'face_{face_count}.jpg')
                cv2.imwrite(face_filename, crop_img)

                label = testet.testimage(face_filename)

                if os.path.exists(face_filename):
                    os.remove(face_filename)

                color = (0, 0, 255) if label == 'fake' else (0, 255, 0)
                cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
                cv2.putText(frame, label, (x, y + 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
    return opencv_to_pil(frame)