import os
import logging

import argparse

import cv2

from protestswap.protestswap import ProtestFaceSwapper

logger = logging.getLogger("protestswap")

DEFAULT_INSIGHTFACE_DIR = "/var/lib/insightface"
INSIGHTFACE_VAR = "INSIGHTFACE_ROOT_DIR"


def main():
    parser = argparse.ArgumentParser(
        prog="protestswap", description="Swaps faces of people on protests"
    )
    parser.add_argument("target", help="The image to swap the faces on")
    parser.add_argument(
        "-t", "--template", action="append", help="The template faces to paste on"
    )

    args = parser.parse_args()

    print(args.target)

    if not args.template:
        print("You need at least one template (with '-t')!")
        exit(1)

    insightface_dir = None
    if os.environ.get(INSIGHTFACE_VAR):
        insightface_dir = os.environ.get(INSIGHTFACE_VAR)
    elif os.path.exists(DEFAULT_INSIGHTFACE_DIR):
        insightface_dir = DEFAULT_INSIGHTFACE_DIR
    else:
        logger.warning(
            f"No directory at '{DEFAULT_INSIGHTFACE_DIR}' and '{INSIGHTFACE_VAR}' not set yet."
        )
        insightface_dir = f"{os.environ.get('HOME')}/.cache/insightface"

    swapper = ProtestFaceSwapper(insightface_dir)
    swapper.prepare_model()

    swapper.set_target_path(args.target)

    for template in args.template:
        print(f"Adding template {template}...")
        swapper.add_template_path(template)

    print("Detecting faces")
    swapper.detect_target_faces()

    swapped = swapper.swap()
    cv2.imwrite("result.png", swapped)


if __name__ == "__main__":
    main()
