import numpy as np
import wave
from scipy.signal import find_peaks
from scipy.fft import fft, fftfreq
import gradio as gr
import tempfile
from PIL import Image
import io

def log_message(message):
    print(message)

def smooth_signal(signal, window_len=11):
    """Apply a smoothing filter to the signal to reduce noise."""
    if window_len < 3:
        return signal
    s = np.r_[signal[window_len-1:0:-1], signal, signal[-2:-window_len-1:-1]]
    window = np.hanning(window_len)
    smoothed = np.convolve(window/window.sum(), s, mode='valid')
    return smoothed

def detect_beeps(audio_data, framerate, freq_zero, freq_one, window_size=2048, step_size=512):
    log_message("Starting beep detection with sliding window.")
    
    binary_data = bytearray()
    num_windows = len(audio_data) // step_size
    previous_bit = None
    
    for i in range(0, len(audio_data) - window_size, step_size):
        # Get the window of data
        window = audio_data[i:i + window_size]
        
        # Perform FFT to detect the dominant frequency in the window
        yf = fft(window)
        xf = fftfreq(window_size, 1 / framerate)
        magnitude = np.abs(yf[:window_size // 2])
        dominant_freq = xf[np.argmax(magnitude)]
        
        # Classify based on the dominant frequency
        if np.isclose(dominant_freq, freq_zero, atol=10):
            current_bit = 0
        elif np.isclose(dominant_freq, freq_one, atol=10):
            current_bit = 1
        else:
            continue  # Skip if it doesn't match expected frequencies
        
        # Only record the bit if it differs from the previous bit (avoids repetition)
        if current_bit != previous_bit:
            binary_data.append(current_bit)
            log_message(f"Detected bit: {current_bit} (Frequency: {dominant_freq} Hz)")
            previous_bit = current_bit
    
    log_message("Finished beep detection.")
    return binary_data

def audio_to_binary(audio_file, freq_zero, freq_one):
    log_message("Starting audio_to_binary function.")
    
    # Read the audio file
    with wave.open(audio_file.name, 'rb') as wav_file:
        framerate = wav_file.getframerate()
        n_frames = wav_file.getnframes()
        audio_data = wav_file.readframes(n_frames)
    
    audio_data = np.frombuffer(audio_data, dtype=np.int16).astype(np.float32)
    audio_data = audio_data / np.max(np.abs(audio_data))  # Normalize to [-1, 1]
    
    # Detect beeps and reconstruct binary data
    binary_data = detect_beeps(audio_data, framerate, freq_zero, freq_one)
    
    if not binary_data:
        return None, "No binary data detected."
    
    # Convert binary list to bytes
    bit_accumulator = 0
    bit_count = 0
    final_binary_data = bytearray()
    
    for bit in binary_data:
        bit_accumulator = (bit_accumulator << 1) | bit
        bit_count += 1
        if bit_count == 8:
            final_binary_data.append(bit_accumulator)
            bit_count = 0
            bit_accumulator = 0
    
    if bit_count > 0:
        final_binary_data.append(bit_accumulator << (8 - bit_count))
    
    log_message(f"Extracted binary data: {final_binary_data[:20]}... (truncated for log)")
    return final_binary_data

def binary_to_file(binary_data, file_format):
    log_message("Starting binary_to_file function.")
    
    temp_file_path = None
    if file_format == 'PNG':
        try:
            image = Image.open(io.BytesIO(binary_data))
            with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_file:
                image.save(temp_file, format='PNG')
                temp_file_path = temp_file.name
            log_message("Successfully reconstructed PNG file.")
        except Exception as e:
            log_message(f"Failed to reconstruct PNG file: {e}")
    elif file_format == 'TXT':
        try:
            text_data = binary_data.decode('utf-8', errors='replace')
            with tempfile.NamedTemporaryFile(delete=False, suffix='.txt') as temp_file:
                temp_file.write(text_data.encode('utf-8'))
                temp_file_path = temp_file.name
            log_message("Successfully reconstructed TXT file.")
        except Exception as e:
            log_message(f"Failed to reconstruct TXT file: {e}")
    else:
        log_message("Unsupported file format.")
    
    return temp_file_path

def process_audio(file, freq_zero, freq_one, file_format):
    log_message("Starting process_audio function.")
    binary_data = audio_to_binary(file, freq_zero, freq_one)
    
    if not binary_data:
        log_message("No binary data was extracted from the audio.")
        return None, "No data was extracted from the audio file."
    
    original_file_path = binary_to_file(binary_data, file_format)
    
    if original_file_path is None:
        return None, "Error in reconstructing the file."
    
    log_message("Processing complete.")
    return original_file_path, "File successfully reconstructed."

gr.Interface(
    fn=process_audio,
    inputs=[
        gr.File(label="Upload Binary WAV File"),
        gr.Number(label="Frequency for 0 (Hz)", value=1000),
        gr.Number(label="Frequency for 1 (Hz)", value=2000),
        gr.Dropdown(choices=['PNG', 'TXT'], label="File Format")
    ],
    outputs=[gr.File(label="Download Original File"), gr.Text(label="Status")],
    title="Binary WAV to Original File Decoder",
    description="Detect beeps with different frequencies in a binary WAV file to reconstruct the original file."
).launch()