import os
import tempfile
import gradio as gr
from huggingface_hub import snapshot_download

# Initialize and import DeepFilterNet components
try:
    from df.enhance import init_df, enhance, load_audio, save_audio
except ImportError:
    raise ImportError("Failed to load DeepFilterNet modules. Make sure it is installed via requirements.txt")

# Configuration
REPO_ID = "detail-co/clear"

print("Downloading model checkpoints from Hugging Face Hub...")
# Downloads the entire target folder structure conforming to DeepFilterNet rules
model_dir = snapshot_download(repo_id=REPO_ID)
print(f"Model successfully saved to local directory: {model_dir}")

# Initialize DeepFilterNet using the downloaded fine-tuned weights
try:
    model, df_state, _ = init_df(model_base_dir=model_dir)
    print("Successfully initialized detail-co/clear model configuration.")
except Exception as e:
    print(f"Standard init failed: {e}. Trying alternative fallback...")
    model, df_state, _ = init_df(model_base_dir=model_dir, epoch=None)

def process_audio(audio_filepath):
    if audio_filepath is None:
        return None
    
    try:
        # Extract native model target sampling rate (usually 48000Hz)
        target_sr = df_state.sr()
        
        # Load and automatically match sample-rate to model expectations
        audio_tensor, _ = load_audio(audio_filepath, sr=target_sr)
        
        # Execute background noise & reverb extraction
        enhanced_tensor = enhance(model, df_state, audio_tensor)
        
        # Save output to a safe, dynamically allocated temporary wave file
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_wav:
            output_filepath = temp_wav.name
            
        save_audio(output_filepath, enhanced_tensor, target_sr)
        return output_filepath
        
    except Exception as error:
        raise gr.Error(f"An error occurred while processing the audio: {str(error)}")

# UI Styling and Layout Design
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="indigo")) as demo:
    gr.Markdown(
        """
        # 🎙️ Clear Voice Enhancer (`detail-co/clear`)
        Remove background noise, hums, and heavy room echo instantly using on-device full-band 48 kHz deep filtering.
        
        **Instructions:** Upload your audio file below, click **Enhance Audio**, then listen to the preview or use the top-right button on the player to download your clean file.
        """
    )
    
    with gr.Row():
        with gr.Column():
            input_audio = gr.Audio(
                label="Upload Raw/Noisy Audio", 
                type="filepath"
            )
            submit_btn = gr.Button("⚡ Enhance Audio", variant="primary")
            
        with gr.Column():
            output_audio = gr.Audio(
                label="Processed Audio (Preview / Download)", 
                type="filepath",
                interactive=False
            )
            
    # Trigger workflow
    submit_btn.click(
        fn=process_audio, 
        inputs=input_audio, 
        outputs=output_audio
    )

    gr.Markdown(
        """
        ---
        *Built using DeepFilterNet 3 fine-tunes optimized for delivering clean, podcast-ready voice tracks.*
        """
    )

# ... keep all your existing app.py code the same above this line ...

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)