import gradio as gr
import fastapi
import starlette
import pydantic
import huggingface_hub

print("GRADIO =", gr.__version__)
print("FASTAPI =", fastapi.__version__)
print("STARLETTE =", starlette.__version__)
print("PYDANTIC =", pydantic.__version__)
print("HF HUB =", huggingface_hub.__version__)

import subprocess
import os
from PIL import Image
import numpy as np
from pydub import AudioSegment


# ----------------------------
# Save audio (numpy -> mp3)
# ----------------------------
def save_audio_mp3(audio_tuple, filename):
    sampling_rate, audio_data = audio_tuple

    audio_bytes = np.array(audio_data, dtype=np.int16).tobytes()

    audio_segment = AudioSegment(
        audio_bytes,
        sample_width=2,
        frame_rate=sampling_rate,
        channels=1
    )

    audio_segment.export(filename, format="mp3")


# ----------------------------
# Merge video + audio (ffmpeg)
# ----------------------------
def merge_audio_video(video_path, audio_path, output_path):

    if os.path.exists(output_path):
        os.remove(output_path)

    cmd = [
        "ffmpeg",
        "-y",
        "-i", video_path,
        "-i", audio_path,
        "-c:v", "copy",
        "-c:a", "aac",
        "-map", "0:v:0",
        "-map", "1:a:0",
        output_path
    ]

    subprocess.run(cmd, check=True)

    return output_path


# ----------------------------
# Inference function
# ----------------------------
def run_inference(input_image, input_audio):

    if input_image is None:
        raise gr.Error("Please upload an image.")

    if input_audio is None:
        raise gr.Error("Please upload audio.")

    os.makedirs("sample_data", exist_ok=True)
    os.makedirs("results", exist_ok=True)

    # Save image
    image_path = "sample_data/uploaded_image.png"
    Image.fromarray(input_image.astype(np.uint8)).save(image_path)

    # Save audio
    audio_path = "sample_data/uploaded_audio.mp3"
    save_audio_mp3(input_audio, audio_path)

    # Run Wav2Lip
    cmd = [
        "python3",
        "inference.py",
        "--checkpoint_path", "checkpoints/wav2lip_gan.pth",
        "--face", image_path,
        "--audio", audio_path
    ]

    result = subprocess.run(
        cmd,
        capture_output=True,
        text=True
    )

    if result.returncode != 0:

        # نجمع stdout و stderr لأن بعض الرسائل قد تظهر في أي منهما
        error = (result.stderr or "") + (result.stdout or "")

        # رسالة عدم اكتشاف الوجه
        if "Face not detected!" in error or "No face detected" in error:
            raise gr.Error(
                "❌ No face detected. Please upload a clear front-facing image."
            )

        # رسالة الصوت غير الصالح
        if "Mel contains nan" in error:
            raise gr.Error(
                "❌ Invalid audio file. Please upload another audio."
            )

        # أي خطأ آخر
        raise gr.Error("❌ Failed to generate video.")

    wav2lip_video = "results/result_voice.mp4"

    if not os.path.exists(wav2lip_video):
        raise gr.Error("Wav2Lip output not found!")

    # merge audio + video
    final_video = merge_audio_video(
        wav2lip_video,
        audio_path,
        "results/final_output.mp4"
    )

    return final_video


# ----------------------------
# UI
# ----------------------------
def create_demo():

    with gr.Blocks() as demo:

        gr.Markdown("# 🎤 Wav2Lip Demo")

        with gr.Row():
            input_image = gr.Image(
                type="numpy",
                label="Input Image"
            )

            input_audio = gr.Audio(
                type="numpy",
                label="Input Audio"
            )

            output_video = gr.Video(
                label="Output Video"
            )

        btn = gr.Button("Generate Video")

        btn.click(
            fn=run_inference,
            inputs=[input_image, input_audio],
            outputs=output_video
        )

        gr.Markdown("## Sample")

        with gr.Row():
            gr.Image(
                "sample/spark.png",
                label="Sample Image"
            )

            gr.Audio(
                "sample/spark_1.1.mp3",
                label="Sample Audio"
            )

            gr.Video(
                "sample/final_output.mp4",
                label="Sample Output"
            )

    return demo


if __name__ == "__main__":
    demo = create_demo()
    demo.queue()
    demo.launch(show_api=True)