import os
import random
import numpy as np
import gradio as gr
from PIL import Image
from pydub import AudioSegment

# Function to generate a corrupted file name
def generate_corrupted_name(original_name):
    base, ext = os.path.splitext(original_name)
    corrupted_name = f"{base}_corrupted{random.randint(1000, 9999)}{ext}"
    return corrupted_name

# Function to glitch images with extreme corruption
def glitch_image(image_path):
    img = Image.open(image_path)
    img = img.convert("RGB")  # Ensure image is in RGB mode
    img_array = np.array(img)

    # Apply extreme corruption by randomly selecting a large number of pixels
    for _ in range(100000):  # Increased number of corrupted pixels
        x = random.randint(0, img_array.shape[1] - 1)
        y = random.randint(0, img_array.shape[0] - 1)
        img_array[y, x] = [random.randint(0, 255) for _ in range(3)]  # Corrupt pixel

    # Additional extreme effects
    if random.random() > 0.5:
        img_array = np.flipud(img_array)  # Flip upside down
    if random.random() > 0.5:
        img_array = np.fliplr(img_array)  # Flip left-right

    if random.random() > 0.5:
        img = Image.fromarray(img_array)
        angle = random.randint(0, 360)
        img = img.rotate(angle)
        img_array = np.array(img)  # Update array after rotation

    # Create a new image from the corrupted array
    corrupted_image = Image.fromarray(img_array)
    
    corrupted_image_name = generate_corrupted_name(image_path)
    corrupted_image.save(corrupted_image_name)
    return corrupted_image_name

# Function to glitch audio
def glitch_audio(audio_path):
    try:
        audio = AudioSegment.from_file(audio_path)
    except Exception as e:
        return f"Error loading audio: {str(e)}"
    
    # Randomly manipulate audio segments for extreme corruption
    for _ in range(15):  # Increased manipulation for more extreme corruption
        start_time = random.randint(0, len(audio) - 1000)
        end_time = start_time + random.randint(500, 3000)
        audio = audio[:start_time] + audio[start_time:end_time].reverse() + audio[end_time:]

    corrupted_audio_name = generate_corrupted_name(audio_path)
    audio.export(corrupted_audio_name, format="wav")  # Export as WAV for simplicity
    return corrupted_audio_name

# Function to glitch video and combine with corrupted audio
def glitch_video(video_path):
    audio_extracted = "extracted_audio.wav"
    extraction_command = f"ffmpeg -y -i {video_path} -q:a 0 -map a {audio_extracted}"
    
    # Run extraction command
    os.system(extraction_command)

    # Check if the audio extraction was successful
    if not os.path.exists(audio_extracted):
        return "Audio extraction failed. Please check the video file."

    # Glitch the extracted audio
    corrupted_audio_path = glitch_audio(audio_extracted)

    # Check if the audio was corrupted successfully
    if "Error" in corrupted_audio_path:
        return corrupted_audio_path

    # Corrupt the video using a filter for visual glitching
    corrupted_video_path = generate_corrupted_name(video_path)
    os.system(f"ffmpeg -y -i {video_path} -vf 'scale=iw*random(0.5)+iw*0.5:ih*random(0.5)+ih*0.5' {corrupted_video_path}")

    # Combine the corrupted video and audio
    final_output_path = generate_corrupted_name(video_path)
    os.system(f"ffmpeg -y -i {corrupted_video_path} -i {corrupted_audio_path} -c:v copy -c:a aac {final_output_path}")

    # Clean up temporary files
    for temp_file in [audio_extracted, corrupted_video_path, corrupted_audio_path]:
        if os.path.exists(temp_file):
            os.remove(temp_file)

    return final_output_path

# Function to glitch text files
def glitch_text(text_path):
    with open(text_path, 'r') as file:
        text_content = file.read()

    # Randomly corrupt text by replacing random characters
    text_list = list(text_content)
    for _ in range(50):  # Increase for more corruption
        index = random.randint(0, len(text_list) - 1)
        text_list[index] = random.choice(['@', '#', '!', '%', '^', '&', '*', '(', ')'])

    corrupted_text = ''.join(text_list)

    corrupted_text_name = generate_corrupted_name(text_path)
    with open(corrupted_text_name, 'w') as file:
        file.write(corrupted_text)

    return corrupted_text_name

# Gradio interface
def glitch_file(file):
    file_extension = file.name.split('.')[-1].lower()
    
    if file_extension in ['png', 'jpg', 'jpeg']:
        return glitch_image(file.name)
    elif file_extension in ['mp4', 'mov', 'avi']:
        return glitch_video(file.name)
    elif file_extension in ['wav', 'mp3']:
        return glitch_audio(file.name)
    elif file_extension in ['txt']:
        return glitch_text(file.name)
    else:
        return "Unsupported file format."

# Gradio UI setup
with gr.Blocks() as app:
    gr.Markdown("### Glitch Any File")
    file_input = gr.File(label="Upload a file (image, video, audio, or text)")
    output = gr.File(label="Corrupted File")
    submit_btn = gr.Button("Glitch")
    submit_btn.click(glitch_file, inputs=file_input, outputs=output)

app.launch()