import cv2
from video_utils import extract_frames,recombine_frames
from face_detection import (detect_faces_insightface_from_frames,
                            extract_faces_from_frames_folder)
from grouping import group_faces_by_identity_facenet_fixed,split_csv_by_identity
from faceswap import face_swap_on_frame_folder,face_swap_multiple_identities
from generate import generate_face_image
import os
from load_models import warm_up_models
from csv_utils import temporal_fill_gaps_in_csv,temporal_smooth_csv

def anonymise_video(input_path: str, output_path: str = None) -> None:
    """
    Anonymises a video by detecting faces, clustering identities, generating synthetic replacements,
    swapping them, and recombining the frames into a new video.
    Args:
        input_path (str): Path to the input video.
        output_path (str): Optional path to save the anonymised video.
    """
    print(f"Anonymising video: {input_path}")

    frames_folder = 'frames'
    csv_path = 'meta_data/frames_detections.csv'
    output_faces = 'output_faces'
    num_identities = 2
    output_folder = 'output_frames'

    warm_up_models()
    '''
    # Step 1: Extract frames from the video
    #extract_frames(input_path, frames_folder)
    # Step 2: Detect faces and save to CSV
    #detect_faces_insightface_from_frames(frames_folder, csv_path)
    # Step 3: Crop detected faces
    #extract_faces_from_frames_folder(frames_folder, csv_path=csv_path, padding=30, output_folder=output_faces)
    # Step 4: Group faces into identities
    group_faces_by_identity_facenet_fixed(faces_folder=output_faces, num_identities=num_identities)
    # Step 5: Split detection CSV into per-identity CSVs
    split_csv_by_identity(original_csv=csv_path)
    '''

    # Step 6: Generate synthetic faces
    generated_faces = []
    for i in range(num_identities):
        print(f'🧬 Generating face for identity: {i}')
        generated_faces.append(generate_face_image())

    # Step 7: Fix frames with no faces
    csv_identities_dir = 'meta_data/identity_csvs'
    identity_csv_paths = []
    for i in range(num_identities):
        path = os.path.join(csv_identities_dir, f"identity_{i}.csv")
        if not os.path.exists(path):
            raise FileNotFoundError(f"Expected CSV not found: {path}")
        identity_csv_paths.append(path)

    identity_csv_paths = fix_faces(num_identities,identity_csv_paths)

    # Step 8: Swap each identity into frames
    face_swap_multiple_identities(
        frames_folder,
        output_folder = output_folder,
        generated_images=generated_faces,
        identity_csv_paths=identity_csv_paths
    )

    # Step 9: Recombine frames into anonymised video
    if output_path is None:
        base, ext = os.path.splitext(input_path)
        output_path = f"{base}_anonymised.mp4"


    recombine_frames(input_path,frames_folder = output_folder,output_video=output_path)


def fix_faces(num_identities,identity_csv_paths):
    new_csv_paths = []
    for i in range(num_identities):
        csv_path = identity_csv_paths[i]
        csv_path = temporal_fill_gaps_in_csv(csv_path)
        #csv_path = temporal_smooth_csv(csv_path)
        new_csv_paths.append(csv_path)

    return new_csv_paths

if __name__ == "__main__":
    input_video = "/Users/sophiemaw/Downloads/CONFIDENTIAL DO NOT SHARE Edna & Paul 29.10.10 Part 2 00.12.46.531.mov"
    output_video = 'anonymised.mp4'
    anonymise_video(input_video)