import cv2
import numpy as np
from PIL import Image
import os
import gradio as gr
from huggingface_hub import hf_hub_download, snapshot_download
import shlex
import subprocess

subprocess.run(shlex.split('pip install flash-attn --no-build-isolation'),
               env=os.environ | {'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"})

def resolve_hf_path(path):
    if isinstance(path, str) and path.startswith("hf://"):
        parts = path[len("hf://"):].split("@")
        if len(parts) == 1:
            repo_id = parts[0]
            filename = None
        elif len(parts) == 2:
            repo_id, filename = parts
        else:
            raise ValueError(f"Invalid HF URI format: {path}")
        token = os.environ.get("HUGGINGFACE_HUB_TOKEN")
        if token is None:
            raise ValueError("HUGGINGFACE_HUB_TOKEN environment variable not set!")
        # If filename is provided, download that file; otherwise, download the whole repo snapshot.
        local_path = hf_hub_download(repo_id=repo_id, filename=filename, token=token) if filename else snapshot_download(repo_id=repo_id, token=token)
        return local_path
    return path

os.environ["FLUX_FILL_PATH"] = "hf://black-forest-labs/FLUX.1-Fill-dev"
os.environ["PORTRAIT_MODEL_PATH"] = "ms://iic/ACE_Plus@portrait/comfyui_portrait_lora64.safetensors"
os.environ["SUBJECT_MODEL_PATH"] = "ms://iic/ACE_Plus@subject/comfyui_subject_lora16.safetensors"
os.environ["LOCAL_MODEL_PATH"] = "ms://iic/ACE_Plus@local_editing/comfyui_local_lora16.safetensors"
os.environ["ACE_PLUS_FFT_MODEL"] = "hf://ali-vilab/ACE_Plus@ace_plus_fft.safetensors"

flux_full = resolve_hf_path(os.environ["FLUX_FILL_PATH"])
ace_plus_fft_model_path = resolve_hf_path(os.environ["ACE_PLUS_FFT_MODEL"])

# Update the environment variables with the resolved local file paths.
os.environ["ACE_PLUS_FFT_MODEL"] = ace_plus_fft_model_path
os.environ["FLUX_FILL_PATH"] = flux_full

from inference.ace_plus_inference import ACEInference
from scepter.modules.utils.config import Config
from modules.flux import FluxMRModiACEPlus
from inference.registry import INFERENCES


config_path = os.path.join("config", "ace_plus_fft.yaml")
cfg = Config(load=True, cfg_file=config_path)
# Instantiate the ACEInference object.
ace_infer = ACEInference(cfg)

def create_face_mask(pil_image):
    """
    Create a binary mask (PIL Image) from a PIL image by detecting the face region.
    The mask will be white (255) on the detected face area and black (0) elsewhere.
    An ellipse is used to better match the shape of a face.
    """
    try:
        # Convert the PIL image to a numpy array in RGB format
        image_np = np.array(pil_image.convert("RGB"))
        # Convert to grayscale for face detection
        gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)

        # Load the Haar cascade for face detection
        cascade_path = cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
        face_cascade = cv2.CascadeClassifier(cascade_path)

        # Detect faces in the image
        faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5)

        # Create an empty mask with the same dimensions as the grayscale image
        mask = np.zeros_like(gray, dtype=np.uint8)

        # For each detected face, draw an ellipse instead of a rectangle
        for (x, y, w, h) in faces:
            # Optionally expand the bounding box slightly for a better fit
            padding = 0.2
            x1 = max(0, int(x - w * padding))
            y1 = max(0, int(y - h * padding))
            x2 = min(gray.shape[1], int(x + w * (1 + padding)))
            y2 = min(gray.shape[0], int(y + h * (1 + padding)))

            # Calculate the center and axes for the ellipse
            center = (int((x1 + x2) / 2), int((y1 + y2) / 2))
            axes = (int((x2 - x1) / 2), int((y2 - y1) / 2))
            # Draw a filled ellipse (white) on the mask
            cv2.ellipse(mask, center, axes, 0, 0, 360, 255, -1)

        return Image.fromarray(mask)
    except Exception as e:
        print(f"Error: {e}")
        raise ValueError('A very specific bad thing happened.')

def face_swap_app(target_img, face_img):
    if target_img is None or face_img is None:
        raise ValueError("Both a target image and a face image must be provided.")

    # (Optional) Ensure images are in RGB
    target_img = target_img.convert("RGB")
    face_img = face_img.convert("RGB")

    edit_mask = create_face_mask(face_img)

    output_img, edit_image, change_image, mask, seed = ace_infer(
        reference_image=target_img,
        edit_image=face_img,
        edit_mask=edit_mask,
        prompt="retain face",
        output_height=1024,
        output_width=1024,
        sampler='flow_euler',
        sample_steps=28,
        guide_scale=50,
        repainting_scale=1.0,
        use_change=True,
        keep_pixels=True,
        keep_pixels_rate=0.8,
        seed=-1
    )
    return output_img

# Create the Gradio interface.
iface = gr.Interface(
    fn=face_swap_app,
    inputs=[
        gr.Image(type="pil", label="Target Image"),
        gr.Image(type="pil", label="Face Image")
    ],
    outputs=gr.Image(type="pil", label="Swapped Face Output"),
    title="ACE++ Face Swap Demo",
    description="Upload a target image and a face image to swap the face using the ACE++ model."
)

if __name__ == "__main__":
    iface.launch()