# refacer_bulk.py
#
# Example usage:
# python refacer_bulk.py --input_path ./input --dest_face myface.jpg --facetoreplace face1.jpg --threshold 0.3
#
# Or, to disable similarity check (i.e., just apply the destination face to all detected faces):
# python refacer_bulk.py --input_path ./input --dest_face myface.jpg

import argparse
import os
import time

import cv2
from PIL import Image

from refacer_no_path import Refacer as NoPathRefacer
import pyfiglet

from utils.minio_client import oss_get_image, minio_client, oss_upload_image


def main():
    input_path = [
        "lanecarford/original_image/7450d6e8-bc54-4c85-940c-4a31c879e02f-0-89.png",
        "lanecarford-outfits/outfits/outfit_6420.jpg",
        "lanecarford-outfits/outfits/outfit_7579.jpg"
    ]
    dest_face = "lanecarford/input_face/leijun.jpg"
    facetoreplace = ""
    threshold = 0.2
    force_cpu = False
    colab_performance = False
    input_dir = input_path

    refacer = NoPathRefacer(force_cpu=force_cpu, colab_performance=colab_performance)

    # Load destination and origin face
    dest_img = oss_get_image(oss_client=minio_client, path=dest_face, data_type="cv2")
    if dest_img is None:
        raise ValueError(f"Destination face image not found: {dest_face}")

    origin_img = None
    if facetoreplace:
        origin_img = oss_get_image(oss_client=minio_client, path=facetoreplace, data_type="cv2")
        if origin_img is None:
            raise ValueError(f"Face to replace image not found: {facetoreplace}")

    disable_similarity = origin_img is None

    faces_config = [{
        'origin': origin_img,
        'destination': dest_img,
        'threshold': threshold
    }]

    refacer.prepare_faces(faces_config, disable_similarity=disable_similarity)

    print(f"Processing images from: {input_dir}")
    image_files = list(input_dir)
    supported_exts = {'jpg', 'jpeg', 'png', 'bmp', 'webp'}

    refaced_images_url = []
    for i, image_path in enumerate(image_files):
        ext = image_path.rsplit(".", 1)[1].lower()
        if ext not in supported_exts:
            print(f"Skipping non-image file: {image_path}")
            continue

        print(f"Refacing: {image_path}")
        try:
            refaced_image = refacer.reface_image(str(image_path), faces_config, disable_similarity=disable_similarity)
            refaced_image_rgb = cv2.cvtColor(refaced_image, cv2.COLOR_RGB2BGR)
            image_bytes = cv2.imencode('.jpg', refaced_image_rgb)[1].tobytes()
            req = oss_upload_image(oss_client=minio_client, bucket="lanecarford", object_name=f"refaced_image/refaced{time.time()}.{ext}", image_bytes=image_bytes)
            refaced_images_url.append(f"{req.bucket_name}/{req.object_name}")
            print(f"Saved -> {req.bucket_name}/{req.object_name}")
        except Exception as e:
            print(f"Failed to process {image_path}: {e}")


if __name__ == "__main__":
    main()