import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
from moviepy.editor import VideoFileClip, AudioFileClip
import librosa
import librosa.display
import soundfile as sf
import gradio as gr
import tempfile

# Function for displaying progress
def display_progress(percent, message, progress=gr.Progress()):
    progress(percent, desc=message)

# Function for extracting audio from video
def extract_audio(video_path, progress):
    display_progress(0.1, "Extracting audio from video", progress)
    try:
        video = VideoFileClip(video_path)
        if video.audio is None:
            raise ValueError("No audio found in the video")
        audio_path = "extracted_audio.wav"
        video.audio.write_audiofile(audio_path)
        display_progress(0.2, "Audio extracted", progress)
        return audio_path
    except Exception as e:
        display_progress(0.2, f"Failed to extract audio: {e}", progress)
        return None

# Function for dividing video into frames
def extract_frames(video_path, progress):
    display_progress(0.3, "Extracting frames from video", progress)
    try:
        video = cv2.VideoCapture(video_path)
        frames = []
        success, frame = video.read()
        while success:
            frames.append(frame)
            success, frame = video.read()
        video.release()
        display_progress(0.4, "Frames extracted", progress)
        return frames
    except Exception as e:
        display_progress(0.4, f"Failed to extract frames: {e}", progress)
        return None

# Convert frame to spectrogram
def frame_to_spectrogram(frame, sr=22050):
    gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    S = np.flipud(gray_frame.astype(np.float32) / 255.0 * 100.0)
    y = librosa.griffinlim(S)
    return y

# Saving audio
def save_audio(y, sr=22050):
    audio_path = 'output_frame_audio.wav'
    sf.write(audio_path, y, sr)
    return audio_path

# Saving frame spectrogram
def save_spectrogram_image(S, frame_number, temp_dir):
    plt.figure(figsize=(10, 4))
    librosa.display.specshow(S)
    plt.tight_layout()
    image_path = os.path.join(temp_dir, f'spectrogram_frame_{frame_number}.png')
    plt.savefig(image_path)
    plt.close()
    return image_path

# Processing all video frames
def process_video_frames(frames, sr=22050, temp_dir=None, progress=gr.Progress()):
    processed_frames = []
    total_frames = len(frames)
    for i, frame in enumerate(frames):
        y = frame_to_spectrogram(frame, sr)
        S = librosa.feature.melspectrogram(y=y, sr=sr)
        image_path = save_spectrogram_image(S, i, temp_dir)
        processed_frame = cv2.imread(image_path)
        processed_frames.append(processed_frame)
        display_progress(0.5 + int((i + 1) / total_frames * 0.7), f"Frame processing {i + 1}/{total_frames}", progress)
    display_progress(0.8, "All frames processed", progress)
    return processed_frames

# Saving video from frames
def save_video_from_frames(frames, output_path, fps=30):
    height, width, layers = frames[0].shape
    video = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
    for frame in frames:
        video.write(frame)
    video.release()

# Adding audio back to video
def add_audio_to_video(video_path, audio_path, output_path, progress):
    display_progress(0.9, "Adding audio back to video", progress)
    try:
        video = VideoFileClip(video_path)
        audio = AudioFileClip(audio_path)
        final_video = video.set_audio(audio)
        final_video.write_videofile(output_path, codec='libx264', audio_codec='aac')
        display_progress(1, "Video's ready", progress)
    except Exception as e:
        display_progress(1, f"Failed to add audio to video: {e}", progress)

# Gradio interface
def process_video(video_path, progress=gr.Progress()):
    try:
        video = VideoFileClip(video_path)
        if video.duration > 10:
            video = video.subclip(0, 10)
            temp_trimmed_video_path = "trimmed_video.mp4"
            video.write_videofile(temp_trimmed_video_path, codec='libx264')
            video_path = temp_trimmed_video_path
    except Exception as e:
        return f"Failed to load video: {e}"

    audio_path = extract_audio(video_path, progress)
    if audio_path is None:
        return "Failed to extract audio from video."
    frames = extract_frames(video_path, progress)
    if frames is None:
        return "Failed to extract frames from video."

    # Creating a temporary folder for saving frames
    with tempfile.TemporaryDirectory() as temp_dir:
        processed_frames = process_video_frames(frames, temp_dir=temp_dir, progress=progress)
        temp_video_path = os.path.join(temp_dir, 'processed_video.mp4')
        save_video_from_frames(processed_frames, temp_video_path)
        output_video_path = 'output_video_with_audio.mp4'
        add_audio_to_video(temp_video_path, audio_path, output_video_path, progress)
        return output_video_path

with gr.Blocks(title='Video from Spectrogram', css="footer{display:none !important}", theme=gr.themes.Soft(primary_hue="green", secondary_hue="green", spacing_size="sm", radius_size="lg")) as iface:
    
    with gr.Group():
        with gr.Row(variant='panel'):
            with gr.Column():
                gr.HTML("<center><h2><a href='https://t.me/pol1trees'>Telegram Channel</a></h2></center>")
            with gr.Column():
                gr.HTML("<center><h2><a href='https://t.me/+GMTP7hZqY0E4OGRi'>Telegram Chat</a></h2></center>")
            with gr.Column():
                gr.HTML("<center><h2><a href='https://www.youtube.com/@Politrees?sub_confirmation=1'>YouTube</a></h2></center>")
            with gr.Column():
                gr.HTML("<center><h2><a href='https://github.com/Bebra777228/Audio-Steganography'>GitHub</a></h2></center>")

    with gr.Column():
        video_input = gr.Video(label="Upload video")
    with gr.Column():
        generate_button = gr.Button("Generate")
    with gr.Column():
        video_output = gr.Video(label="VideoSpectrogram")

        def gradio_video_process_fn(video_input, progress=gr.Progress()):
            return process_video(video_input, progress)

        generate_button.click(
            gradio_video_process_fn,
            inputs=[video_input],
            outputs=[video_output]
        )

iface.launch(share=True)