import cv2
from ultralytics import YOLO
from moviepy.editor import VideoFileClip

model = YOLO("yolov8-face-hf.pt")

# ------------------ IMAGE FUNCTION ------------------ #
def blur_faces_image(image):
    results = model(image)

    for result in results:
        for box in result.boxes.xyxy.cpu().numpy():
            x1, y1, x2, y2 = map(int, box)
            x1, y1 = max(0, x1), max(0, y1)
            x2, y2 = min(image.shape[1], x2), min(image.shape[0], y2)

            face = image[y1:y2, x1:x2]
            if face.size > 0:
                blurred = cv2.GaussianBlur(face, (101, 101), 30)
                image[y1:y2, x1:x2] = blurred

    return image

# ------------------ VIDEO FUNCTION ------------------ #
def blur_faces_video(input_path, output_path, update_callback=None):
    cap = cv2.VideoCapture(input_path)
    if not cap.isOpened():
        raise IOError("Cannot open video")

    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    width = width if width % 2 == 0 else width - 1
    height = height if height % 2 == 0 else height - 1

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    temp_video_path = "blurred_temp.mp4"
    out = cv2.VideoWriter(temp_video_path, fourcc, fps, (width, height))

    frame_num = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        frame = cv2.resize(frame, (width, height))
        results = model(frame, verbose=False)[0]
        h, w = frame.shape[:2]

        for box in results.boxes.xyxy:
            x1, y1, x2, y2 = map(int, box)
            x1 = max(0, min(w, x1))
            x2 = max(0, min(w, x2))
            y1 = max(0, min(h, y1))
            y2 = max(0, min(h, y2))

            face = frame[y1:y2, x1:x2]
            if face.size > 0:
                blur_size = max(55, ((x2 - x1) // 2) | 1)
                if blur_size % 2 == 0:
                    blur_size += 1
                blurred = cv2.GaussianBlur(face, (blur_size, blur_size), 0)
                frame[y1:y2, x1:x2] = blurred

        out.write(frame)
        frame_num += 1

        if update_callback and total_frames > 0:
            update_callback(frame_num / total_frames)

    cap.release()
    out.release()

    # Merge audio
    try:
        if update_callback:
            update_callback(0.99)

        original = VideoFileClip(input_path)
        processed = VideoFileClip(temp_video_path).set_audio(original.audio)
        processed.write_videofile(
            output_path,
            codec="libx264",
            audio_codec="aac",
            threads=1,
            logger=None
        )

        if update_callback:
            update_callback(1.0)

        return output_path

    except Exception as e:
        print("Audio merging failed:", e)
        return temp_video_path