import os
import gdown
import insightface
from insightface.app import FaceAnalysis
from src.utils.logger import logger
import warnings

warnings.filterwarnings("ignore", category=FutureWarning)

class FaceSwapper:
    def __init__(self, app_name='buffalo_l', det_size=(640, 640), device="cpu"):
        model_dir = "/tmp/.insightface/models"
        os.makedirs(model_dir, exist_ok=True)
        self.device = device
        self.app = FaceAnalysis(name=app_name, root=model_dir)
        self.app.prepare(ctx_id=0 if device == "cuda" else -1, det_size=det_size)
        self.swapper = None

        logger.info('FaceSwapper initialized')


    def load_swapper_model(self, model_url, model_path):
        # Set up the gdown cache directory
        gdown_cache_dir = "/tmp/.cache/gdown"
        os.makedirs(gdown_cache_dir, exist_ok=True)
        os.environ['XDG_CACHE_HOME'] = "/tmp/.cache"

        # Log the cache directory for debugging
        logger.info(f"Setting XDG_CACHE_HOME to: {os.environ['XDG_CACHE_HOME']}")

        # Ensure cookies.txt exists
        cookie_path = os.path.join(gdown_cache_dir, "cookies.txt")
        if not os.path.exists(cookie_path):
            with open(cookie_path, 'w') as f:
                # Create an empty cookies.txt file
                pass
        logger.info(f"Ensured cookies.txt exists at: {cookie_path}")

        # Check if model path exists, and download if not
        if not os.path.exists(model_path):
            os.makedirs(os.path.dirname(model_path), exist_ok=True)
            gdown.download(model_url, model_path, quiet=False, use_cookies=False)

        self.swapper = insightface.model_zoo.get_model(model_path, download=False, download_zip=False)
        logger.info('Swapper model loaded')

    def face_swap(self, img1, img2, enhance=False):
        face1 = self.app.get(img1)[0]
        face2 = self.app.get(img2)[0]
        img1_ = img1.copy()
        img1_ = self.swapper.get(img1_, face1, face2, paste_back=True)
        logger.info('Face swapped')
        return img1_