import os
import cv2
import numpy as np
from PIL import Image
from huggingface_hub import hf_hub_download
from tqdm import tqdm

_app = None
_swapper = None

def get_swapper_models():
    global _app, _swapper
    if _app is None or _swapper is None:
        try:
            import insightface
            from insightface.app import FaceAnalysis

            print("Initializing InsightFace CPU model (det_size=320)...")
            _app = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
            _app.prepare(ctx_id=-1, det_size=(320, 320))

            # Download inswapper_128.onnx from HF Hub
            hf_token = os.environ.get("HF_TOKEN") or True
            try:
                model_path = hf_hub_download(
                    repo_id="ezioruan/inswapper_128.onnx",
                    filename="inswapper_128.onnx",
                    token=hf_token
                )
            except Exception as dl_err:
                print(f"Primary repo download notice: {dl_err}. Trying fallback...")
                model_path = hf_hub_download(
                    repo_id="Gourieff/ReActor",
                    filename="models/inswapper_128.onnx",
                    repo_type="dataset",
                    token=hf_token
                )
            _swapper = insightface.model_zoo.get_model(model_path, providers=['CPUExecutionProvider'])
            print("✅ InsightFace CPU Swapper loaded successfully!")
        except Exception as e:
            print(f"⚠️ Face Swapper load warning: {e}")
            _app = None
            _swapper = None
    return _app, _swapper

def swap_face_in_frames(
    source_pil_image: Image.Image, 
    frames_np: list, 
    ref_face_image: Image.Image = None, 
    target_gender: str = "Any / All Faces",
    swap_last_n: int = 4,
    progress=None
) -> list:
    """
    Swaps face from ref_face_image (or source_pil_image) into video frames using InsightFace CPU.
    Supports swap_last_n frames (0 = All Frames). If swap_last_n > total_frames, falls back to 2.
    Runs 100% on CPU (0 GPU quota used).
    """
    app_model, swapper_model = get_swapper_models()
    if app_model is None or swapper_model is None:
        print("⚠️ Face Swapper model unavailable. Returning original frames.")
        return frames_np

    try:
        source_img = ref_face_image if ref_face_image is not None else source_pil_image
        if source_img is None:
            return frames_np

        source_bgr = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR)
        source_faces = app_model.get(source_bgr)

        if not source_faces:
            print("⚠️ No face detected in source/reference image. Skipping face swap.")
            return frames_np

        source_faces.sort(key=lambda x: (x.bbox[2]-x.bbox[0]) * (x.bbox[3]-x.bbox[1]), reverse=True)
        source_face = source_faces[0]

        total_all = len(frames_np)
        swap_last_n = int(swap_last_n)

        # Fallback calculation
        if swap_last_n == 0:
            n_swap = total_all
        elif swap_last_n > total_all:
            print(f"Notice: swap_last_n ({swap_last_n}) exceeds total frames ({total_all}). Fallback to 2 frames.")
            n_swap = min(2, total_all)
        else:
            n_swap = swap_last_n

        if n_swap < total_all:
            unchanged_prefix = list(frames_np[:-n_swap])
            target_frames = list(frames_np[-n_swap:])
        else:
            unchanged_prefix = []
            target_frames = list(frames_np)

        swapped_sub = []
        total_sub = len(target_frames)
        print(f"👤 Processing CPU Face Swap on {total_sub} frames (Last N={n_swap}, Gender filter: {target_gender})...")

        for idx, frame in enumerate(tqdm(target_frames, desc="👤 CPU Face Swap")):
            if progress is not None:
                try:
                    progress((idx + 1) / total_sub, desc=f"👤 Swapping Face on Frame {idx+1}/{total_sub} (CPU)...")
                except Exception:
                    pass

            if isinstance(frame, Image.Image):
                frame_uint8 = cv2.cvtColor(np.array(frame), cv2.COLOR_RGB2BGR)
            elif isinstance(frame, np.ndarray):
                frame_uint8 = (frame * 255).astype(np.uint8) if frame.dtype != np.uint8 else frame.copy()
                frame_uint8 = cv2.cvtColor(frame_uint8, cv2.COLOR_RGB2BGR)
            else:
                frame_uint8 = np.array(frame, dtype=np.uint8)
                frame_uint8 = cv2.cvtColor(frame_uint8, cv2.COLOR_RGB2BGR)

            target_bgr = frame_uint8
            target_faces = app_model.get(target_bgr)

            if target_faces:
                res_bgr = target_bgr.copy()
                for target_face in target_faces:
                    gender_val = getattr(target_face, 'gender', None)
                    sex_val = getattr(target_face, 'sex', None)
                    
                    if target_gender == "Female Faces Only":
                        is_female = (gender_val == 0) or (sex_val == 'F')
                        if not is_female:
                            continue
                    elif target_gender == "Male Faces Only":
                        is_male = (gender_val == 1) or (sex_val == 'M')
                        if not is_male:
                            continue

                    res_bgr = swapper_model.get(res_bgr, target_face, source_face, paste_back=True)

                res_rgb = cv2.cvtColor(res_bgr, cv2.COLOR_BGR2RGB)
                if isinstance(frame, np.ndarray) and frame.dtype != np.uint8:
                    swapped_sub.append(res_rgb.astype(np.float32) / 255.0)
                elif isinstance(frame, Image.Image):
                    swapped_sub.append(Image.fromarray(res_rgb))
                else:
                    swapped_sub.append(res_rgb)
            else:
                swapped_sub.append(frame)

        final_result = unchanged_prefix + swapped_sub
        print(f"✅ CPU Face Swap complete ({len(swapped_sub)} frames swapped)!")
        return final_result

    except Exception as e:
        print(f"⚠️ Face Swapper execution error: {e}")
        return frames_np

def map_gender_param(target_gender: str) -> str:
    if not target_gender:
        return "all"
    tg = str(target_gender).lower()
    if "female" in tg or "wanita" in tg or "perempuan" in tg:
        return "female"
    elif "male" in tg or "pria" in tg or "laki" in tg:
        return "male"
    return "all"

def call_sulphur_faceswap_api(source_img: Image.Image, target_img: Image.Image, target_gender: str = "all", enhance_with_gfpgan: bool = True, server_url: str = None) -> Image.Image:
    """
    Calls Sulphur AI API (/api/v1/faceswap) to perform InsightFace Face Swap + GFPGAN Face Restoration.
    """
    import io
    import requests
    import config

    target_url = server_url or config.SULPHUR_API_URL or os.environ.get("SULPHUR_API_URL", "http://localhost:6666")
    if not target_url or not str(target_url).strip():
        return None

    clean_url = str(target_url).strip().rstrip("/")
    endpoint = f"{clean_url}/api/v1/faceswap"

    try:
        source_bytes = io.BytesIO()
        source_img.convert("RGB").save(source_bytes, format="JPEG", quality=95)
        source_bytes.seek(0)

        target_bytes = io.BytesIO()
        target_img.convert("RGB").save(target_bytes, format="JPEG", quality=95)
        target_bytes.seek(0)

        files = {
            "source_image": ("source.jpg", source_bytes, "image/jpeg"),
            "target_image": ("target.jpg", target_bytes, "image/jpeg")
        }
        data = {
            "enhance_with_gfpgan": "true" if enhance_with_gfpgan else "false",
            "target_gender": map_gender_param(target_gender)
        }

        print(f"🌐 Calling Sulphur AI Face Swap API at {endpoint} (Gender: {map_gender_param(target_gender)})...")
        res = requests.post(endpoint, files=files, data=data, timeout=15)
        if res.status_code == 200 and res.content:
            result_img = Image.open(io.BytesIO(res.content)).convert("RGB")
            print("✅ Sulphur AI Face Swap + GFPGAN API succeeded!")
            return result_img
        else:
            print(f"⚠️ Sulphur AI Face Swap API returned status {res.status_code}")
    except Exception as e:
        print(f"⚠️ Sulphur AI Face Swap API notice: {e}")
    return None

def call_sulphur_enhance_face_api(image: Image.Image, server_url: str = None) -> Image.Image:
    """
    Calls Sulphur AI API (/api/v1/enhance-face) to sharpen & restore face details via GFPGAN v1.4.
    """
    import io
    import requests
    import config

    target_url = server_url or config.SULPHUR_API_URL or os.environ.get("SULPHUR_API_URL", "http://localhost:6666")
    if not target_url or not str(target_url).strip():
        return None

    clean_url = str(target_url).strip().rstrip("/")
    endpoint = f"{clean_url}/api/v1/enhance-face"

    try:
        img_bytes = io.BytesIO()
        image.convert("RGB").save(img_bytes, format="JPEG", quality=95)
        img_bytes.seek(0)

        files = {"image": ("face.jpg", img_bytes, "image/jpeg")}

        print(f"🌐 Calling Sulphur AI GFPGAN Face Enhance API at {endpoint}...")
        res = requests.post(endpoint, files=files, timeout=15)
        if res.status_code == 200 and res.content:
            result_img = Image.open(io.BytesIO(res.content)).convert("RGB")
            print("✅ Sulphur AI GFPGAN Face Enhance succeeded!")
            return result_img
        else:
            print(f"⚠️ Sulphur AI Face Enhance API status {res.status_code}")
    except Exception as e:
        print(f"⚠️ Sulphur AI Face Enhance API notice: {e}")
    return None

def swap_face_in_single_image(
    target_pil_image: Image.Image, 
    ref_face_image: Image.Image = None, 
    target_gender: str = "Any / All Faces",
    enhance_with_gfpgan: bool = True
) -> Image.Image:
    """
    Swaps face on a single PIL image strictly using local CPU InsightFace (0 GPU quota).
    Returns swapped PIL Image.
    """
    if target_pil_image is None:
        return None

    # Execute strictly on local CPU InsightFace
    swapped_frames = swap_face_in_frames(
        source_pil_image=target_pil_image,
        frames_np=[target_pil_image],
        ref_face_image=ref_face_image,
        target_gender=target_gender,
        swap_last_n=0
    )
    res_frame = swapped_frames[0]
    if isinstance(res_frame, Image.Image):
        return res_frame
    elif isinstance(res_frame, np.ndarray):
        if res_frame.dtype != np.uint8:
            res_frame = (res_frame * 255).astype(np.uint8)
        return Image.fromarray(res_frame)
    return target_pil_image