import os
from PIL import Image
import numpy as np
import roop.globals
from roop.core import start, decode_execution_providers, suggest_max_memory, suggest_execution_threads
from roop.processors.frame.core import get_frame_processors_modules

def to_pil(img):
    """Convert np.ndarray to PIL.Image if needed."""
    if isinstance(img, np.ndarray):
        return Image.fromarray(img)
    return img

def swap_face_now(source_img, target_img, do_enhance=True): 
    TEMP_DIR = "temp/"
    os.makedirs(TEMP_DIR, exist_ok=True)

    # ✅ Ép kiểu an toàn và convert về RGB để tránh lỗi JPEG
    source_img = to_pil(source_img)
    target_img = to_pil(target_img)
    
    # Convert về RGB nếu ảnh có alpha channel
    if source_img.mode in ('RGBA', 'LA'):
        # Tạo background trắng và paste ảnh lên
        rgb_source = Image.new('RGB', source_img.size, (255, 255, 255))
        rgb_source.paste(source_img, mask=source_img.split()[-1] if source_img.mode == 'RGBA' else None)
        source_img = rgb_source
    elif source_img.mode != 'RGB':
        source_img = source_img.convert('RGB')
        
    if target_img.mode in ('RGBA', 'LA'):
        # Tạo background trắng và paste ảnh lên
        rgb_target = Image.new('RGB', target_img.size, (255, 255, 255))
        rgb_target.paste(target_img, mask=target_img.split()[-1] if target_img.mode == 'RGBA' else None)
        target_img = rgb_target
    elif target_img.mode != 'RGB':
        target_img = target_img.convert('RGB')

    source_path = os.path.join(TEMP_DIR, "input.jpg")
    target_path = os.path.join(TEMP_DIR, "target.jpg")
    output_path = os.path.join(TEMP_DIR, "output.jpg")

    source_img.save(source_path, 'JPEG')
    target_img.save(target_path, 'JPEG')

    # Roop config
    roop.globals.source_path = source_path
    roop.globals.target_path = target_path
    roop.globals.output_path = output_path
    roop.globals.frame_processors = ["face_swapper", "face_enhancer"] if do_enhance else ["face_swapper"]
    roop.globals.headless = True
    roop.globals.many_faces = False
    roop.globals.max_memory = suggest_max_memory()
    roop.globals.execution_providers = decode_execution_providers(["cuda"])
    roop.globals.execution_threads = suggest_execution_threads()
    roop.globals.keep_fps = False
    roop.globals.keep_audio = False
    roop.globals.keep_frames = False
    roop.globals.video_encoder = None

    for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
        if not frame_processor.pre_check():
            return None

    start()

    if not os.path.exists(output_path):
        return None

    return Image.open(output_path)