from pathlib import Path

import gradio as gr

from pipeline import (
    copy_file_to_dir,
    extract_coordinates,
    extract_face_video,
    make_run_dir,
    merge_lips,
)

BASE_DIR = Path(__file__).resolve().parent
WORK_DIR = BASE_DIR / "work"
WORK_DIR.mkdir(parents=True, exist_ok=True)


def _normalize_upload_path(file_obj):
    if file_obj is None:
        return None
    if isinstance(file_obj, str):
        return file_obj
    return str(file_obj)


def run_step1(original_video):
    try:
        original_path = _normalize_upload_path(original_video)
        if not original_path:
            raise ValueError("Please upload an original video.")

        run_dir = make_run_dir(WORK_DIR, "step1")
        local_video = copy_file_to_dir(original_path, run_dir)

        face_path, lip_path, face_bbox, lip_bbox = extract_coordinates(
            video_path=str(local_video),
            output_dir=str(run_dir),
            face_name="face_coords_avg.pkl",
            lip_name="lip_coords_avg.pkl",
        )

        status = (
            "Step 1 completed. "
            f"Face bbox: {face_bbox}. "
            f"Lip bbox: {lip_bbox}."
        )
        return status, face_path, lip_path
    except Exception as exc:
        return f"Step 1 failed: {exc}", None, None


def run_step2(original_video, face_coords, lip_coords):
    try:
        original_path = _normalize_upload_path(original_video)
        face_path = _normalize_upload_path(face_coords)
        lip_path = _normalize_upload_path(lip_coords)

        if not original_path:
            raise ValueError("Please upload the original video.")
        if not face_path:
            raise ValueError("Please upload face coordinates (.pkl).")

        run_dir = make_run_dir(WORK_DIR, "step2")
        local_video = copy_file_to_dir(original_path, run_dir)
        local_face = copy_file_to_dir(face_path, run_dir, "face_coords_avg.pkl")

        local_lip = None
        if lip_path:
            local_lip = copy_file_to_dir(lip_path, run_dir, "lip_coords_avg.pkl")

        cropped_face_path = run_dir / "cropped_face.mp4"
        extract_face_video(
            video_path=str(local_video),
            face_coords_path=str(local_face),
            output_path=str(cropped_face_path),
        )

        status = "Step 2 completed. Download cropped face video and coordinate files below."
        return status, str(cropped_face_path), str(cropped_face_path), str(local_face), str(local_lip) if local_lip else None
    except Exception as exc:
        return f"Step 2 failed: {exc}", None, None, None, None


def run_step4(original_video, lip_synced_video, face_coords, lip_coords, audio_file):
    try:
        original_path = _normalize_upload_path(original_video)
        lipsynced_path = _normalize_upload_path(lip_synced_video)
        face_path = _normalize_upload_path(face_coords)
        lip_path = _normalize_upload_path(lip_coords)
        audio_path = _normalize_upload_path(audio_file)

        if not original_path:
            raise ValueError("Please upload the original video.")
        if not lipsynced_path:
            raise ValueError("Please upload the lip-synced face video from Step 3.")
        if not face_path:
            raise ValueError("Please upload face coordinates (.pkl).")
        if not lip_path:
            raise ValueError("Please upload lip coordinates (.pkl).")

        run_dir = make_run_dir(WORK_DIR, "step4")
        local_original = copy_file_to_dir(original_path, run_dir, "original_video.mp4")
        local_lipsynced = copy_file_to_dir(lipsynced_path, run_dir, "lip_synced_face_video.mp4")
        local_face = copy_file_to_dir(face_path, run_dir, "face_coords_avg.pkl")
        local_lip = copy_file_to_dir(lip_path, run_dir, "lip_coords_avg.pkl")

        local_audio = None
        if audio_path:
            local_audio = copy_file_to_dir(audio_path, run_dir)

        final_path = run_dir / "final_synced_output.mp4"
        final_video_path, audio_used = merge_lips(
            original_video_path=str(local_original),
            lip_synced_video_path=str(local_lipsynced),
            face_coords_path=str(local_face),
            lip_coords_path=str(local_lip),
            final_output_path=str(final_path),
            audio_path=str(local_audio) if local_audio else None,
        )

        status = f"Step 4 completed. Final video generated. Audio source used: {audio_used}"
        return status, final_video_path, final_video_path
    except Exception as exc:
        return f"Step 4 failed: {exc}", None, None


with gr.Blocks(title="Dub Module - Steps 1, 2, and 4") as demo:
    gr.Markdown(
        """
# Dub Module Gradio App (HF Ready)
This app implements Step 1, Step 2, and Step 4 from your pipeline.
- Step 3 must be done manually outside this app.
- Step 5 is not included.
        """
    )

    with gr.Tab("Step 1 - Extract Coordinates"):
        gr.Markdown("Upload the original video to generate `face_coords_avg.pkl` and `lip_coords_avg.pkl`.")
        s1_video = gr.File(label="Original Video", file_types=["video"], type="filepath")
        s1_run = gr.Button("Run Step 1")
        s1_status = gr.Textbox(label="Status", interactive=False)
        s1_face = gr.File(label="Face Coordinates (.pkl)")
        s1_lip = gr.File(label="Lip Coordinates (.pkl)")
        s1_run.click(fn=run_step1, inputs=[s1_video], outputs=[s1_status, s1_face, s1_lip])

    with gr.Tab("Step 2 - Extract Cropped Face Video"):
        gr.Markdown(
            "Upload original video and face coordinates. Lip coordinates are optional here, "
            "but if provided they are returned for download as requested."
        )
        s2_video = gr.File(label="Original Video", file_types=["video"], type="filepath")
        s2_face = gr.File(label="Face Coordinates (.pkl)", file_types=[".pkl"], type="filepath")
        s2_lip = gr.File(label="Lip Coordinates (.pkl) - optional", file_types=[".pkl"], type="filepath")
        s2_run = gr.Button("Run Step 2")
        s2_status = gr.Textbox(label="Status", interactive=False)
        s2_preview = gr.Video(label="Cropped Face Video Preview")
        s2_video_file = gr.File(label="Download Cropped Face Video")
        s2_face_out = gr.File(label="Download Face Coordinates")
        s2_lip_out = gr.File(label="Download Lip Coordinates")
        s2_run.click(
            fn=run_step2,
            inputs=[s2_video, s2_face, s2_lip],
            outputs=[s2_status, s2_preview, s2_video_file, s2_face_out, s2_lip_out],
        )

    with gr.Tab("Step 3 - Manual (Outside App)"):
        gr.Markdown(
            """
Run your Step 3 lip-sync process manually using the cropped face video from Step 2.
After Step 3, return to Step 4 and upload:
1. Original video
2. Lip-synced face video from your external tool
3. Face coordinates pkl
4. Lip coordinates pkl
5. Optional audio file used during lip-sync
            """
        )

    with gr.Tab("Step 4 - Merge and Final Output"):
        gr.Markdown("Merge the lip-synced lips back to original video and download the final output.")
        s4_original = gr.File(label="Original Video", file_types=["video"], type="filepath")
        s4_lipsynced = gr.File(label="Lip-synced Face Video", file_types=["video"], type="filepath")
        s4_face = gr.File(label="Face Coordinates (.pkl)", file_types=[".pkl"], type="filepath")
        s4_lip = gr.File(label="Lip Coordinates (.pkl)", file_types=[".pkl"], type="filepath")
        s4_audio = gr.File(label="Audio from Step 3 (optional)", file_types=["audio"], type="filepath")
        s4_run = gr.Button("Run Step 4")
        s4_status = gr.Textbox(label="Status", interactive=False)
        s4_preview = gr.Video(label="Final Video Preview")
        s4_file = gr.File(label="Download Final Video")
        s4_run.click(
            fn=run_step4,
            inputs=[s4_original, s4_lipsynced, s4_face, s4_lip, s4_audio],
            outputs=[s4_status, s4_preview, s4_file],
        )


if __name__ == "__main__":
    demo.launch()