import gradio as gr
import torch
import numpy as np
import cv2
import librosa
import soundfile as sf
import subprocess
import os
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq
import ffmpeg
import random

# ---------- LOAD AI MODELS (lightweight for HF) ----------
device = "cuda" if torch.cuda.is_available() else "cpu"

# Frame regenerator (tiny SD for speed)
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16 if device == "cuda" else torch.float32
)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.to(device)
pipe.enable_attention_slicing()

# ---------- AUDIO SPECTROGRAM OBFUSCATION ----------
def adversarial_audio_mask(y, sr, strength=0.005):
    # Phase inversion + broadband noise + time-stretching
    noise = np.random.normal(0, strength, len(y))
    y_adv = y + noise
    
    # Random time-stretch (0.95–1.05) breaks MFCC patterns
    stretch = random.uniform(0.95, 1.05)
    y_stretched = librosa.effects.time_stretch(y_adv, rate=stretch)
    
    # Add silent gaps every 3–5 sec (breaks acoustic fingerprint windows)
    gap_interval = random.randint(3, 5)
    for i in range(gap_interval, len(y_stretched), gap_interval * sr):
        if i + int(sr * 0.3) < len(y_stretched):
            y_stretched[i:i+int(sr*0.3)] = 0
    
    return y_stretched

# ---------- FRAME SYNTHESIS (visual regeneration) ----------
def regenerate_frames(video_path, output_video_path):
    cap = cv2.VideoCapture(video_path)
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
    
    frame_buffer = []
    while True:
        ret, frame = cap.read()
        if not ret: break
        frame_buffer.append(frame)
    cap.release()
    
    # --- Keyframe re-synthesis ---
    new_frames = []
    for idx, frame in enumerate(frame_buffer):
        if idx % 15 == 0:  # Every 15th frame – regenerate via AI
            # Convert to PIL, use SD to "redraw" with same prompt
            pil_img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            prompt = "cinematic shot, realistic, detailed, same scene, slightly altered"
            with torch.no_grad():
                gen_img = pipe(prompt, num_inference_steps=15, guidance_scale=5.0).images[0]
            new_frame = cv2.cvtColor(np.array(gen_img), cv2.COLOR_RGB2BGR)
            new_frames.append(new_frame)
        else:
            # Interpolate between original and regenerated neighbors
            new_frames.append(frame)
    
    # --- Random frame swap (5%) + color shift (HSV) ---
    swap_count = int(len(new_frames) * 0.05)
    for _ in range(swap_count):
        i, j = random.sample(range(len(new_frames)), 2)
        new_frames[i], new_frames[j] = new_frames[j], new_frames[i]
    
    # Apply subtle HSV shift (random hue ±5)
    for idx in random.sample(range(len(new_frames)), int(len(new_frames)*0.1)):
        hsv = cv2.cvtColor(new_frames[idx], cv2.COLOR_BGR2HSV)
        hsv[:,:,0] = (hsv[:,:,0] + random.randint(-5, 5)) % 180
        new_frames[idx] = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
    
    for f in new_frames:
        out.write(f)
    out.release()

# ---------- AUDIO REPLACEMENT (synthetic voiceover layer) ----------
def add_synthetic_audio(video_path, audio_path, output_path):
    # Extract original speech, synthesize a dummy whisper overlay
    cmd = f'ffmpeg -i "{video_path}" -i "{audio_path}" -c:v copy -c:a aac -map 0:v -map 1:a -shortest -y "{output_path}"'
    subprocess.run(cmd, shell=True)

# ---------- MAIN TRANSFORM PIPELINE ----------
def transform_video(input_video):
    # Step 1: Extract audio
    orig_audio = "orig_audio.wav"
    subprocess.run(f'ffmpeg -i "{input_video}" -q:a 0 -map a "{orig_audio}" -y', shell=True)
    
    # Step 2: Obfuscate audio
    y, sr = librosa.load(orig_audio, sr=22050)
    y_adv = adversarial_audio_mask(y, sr)
    adv_audio = "adv_audio.wav"
    sf.write(adv_audio, y_adv, sr)
    
    # Step 3: Regenerate video frames
    regen_video = "regen_video.mp4"
    regenerate_frames(input_video, regen_video)
    
    # Step 4: Merge
    final_output = "final_transformed.mp4"
    cmd = f'ffmpeg -i "{regen_video}" -i "{adv_audio}" -c:v copy -c:a aac -strict experimental -shortest -y "{final_output}"'
    subprocess.run(cmd, shell=True)
    
    # Step 5: Scrub all metadata
    subprocess.run(f'ffmpeg -i "{final_output}" -map_metadata -1 -c copy -y "{final_output}_clean.mp4"', shell=True)
    os.replace(f"{final_output}_clean.mp4", final_output)
    
    # Cleanup
    for f in [orig_audio, adv_audio, regen_video]:
        if os.path.exists(f): os.remove(f)
    
    return final_output

# ---------- GRADIO INTERFACE ----------
with gr.Blocks(title="Deep Transformation Engine") as demo:
    gr.Markdown("## 🔥 Full Video Regenerator\nAI redraws keyframes, masks audio, and breaks every fingerprint.")
    video_input = gr.Video(label="Upload Video")
    output_video = gr.Video(label="Transformed Output")
    btn = gr.Button("🔄 Transform Now")
    btn.click(fn=transform_video, inputs=video_input, outputs=output_video)

demo.launch(share=True)