import os
import torch
import shutil
import uuid
from pathlib import Path
import warnings
from typing import BinaryIO
from src.wav2lip.models import Wav2Lip
from src.wav2lip import inference

warnings.filterwarnings("ignore")

# Directories setup
UPLOAD_DIR = Path("uploads")
OUTPUT_DIR = Path("outputs")
RESULT_TEMP = Path("wav2lip/results/result_voice.mp4")
CHECKPOINT_PATH = "wav2lip_checkpoints/wav2lip_gan.pth"

UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
Path("wav2lip/results").mkdir(parents=True, exist_ok=True)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MODEL = None


def load_model(checkpoint_path: str):
    global MODEL
    if MODEL is not None:
        return MODEL

    print(f"Loading Wav2Lip model from {checkpoint_path} on {DEVICE}")
    model = Wav2Lip()

    # Load checkpoint on the selected device
    checkpoint = torch.load(checkpoint_path, map_location=DEVICE)

    # Fix possible 'module.' prefix in keys if loaded from multi-GPU training
    state_dict = {
        k.replace("module.", ""): v for k, v in checkpoint["state_dict"].items()
    }
    model.load_state_dict(state_dict)

    # Move model to device (GPU or CPU)
    model = model.to(DEVICE)
    model.eval()

    MODEL = model
    return MODEL


def unload_model():
    global MODEL
    if MODEL:
        del MODEL
        MODEL = None
        if DEVICE == "mps":
            torch.mps.empty_cache()
        print("Model unloaded and memory freed.")


def predict(image_file: BinaryIO, audio_file: BinaryIO, model) -> str:
    """
    Receives image/audio files as file-like objects (e.g. BytesIO), saves them, runs prediction,
    and returns the path to the result video.
    """
    # Generate session-specific filenames
    session_id = str(uuid.uuid4())
    image_path = UPLOAD_DIR / f"{session_id}.jpg"
    audio_path = UPLOAD_DIR / f"{session_id}.mp3"
    output_path = OUTPUT_DIR / f"{session_id}_result.mp4"

    # Save input files
    with open(image_path, "wb") as f:
        shutil.copyfileobj(image_file, f)
    with open(audio_path, "wb") as f:
        shutil.copyfileobj(audio_file, f)

    print(f"Saved inputs to: {image_path}, {audio_path}")

    # # Load the model
    # model = load_model(CHECKPOINT_PATH)

    # Run inference
    inference.main(str(image_path), str(audio_path), model)

    # Move output
    if not RESULT_TEMP.exists():
        raise RuntimeError("Inference failed. Output video not found.")

    shutil.move(RESULT_TEMP, output_path)
    print(f"Saved output to: {output_path}")

    return str(output_path)


# from pathlib import Path
# from wav2lip_api import predict, load_model  # Replace with your actual script/module name

# # Load image and audio files as binary
# image_path = "sample_inputs/sample.jpg"  # Replace with actual image path
# audio_path = "sample_inputs/sample.mp3"  # Replace with actual audio path

# with open(image_path, "rb") as img_file, open(audio_path, "rb") as aud_file:
#     result_path = predict(img_file, aud_file)

# print(f"Result video saved at: {result_path}")