"""
app.py — Gradio front-end for tape_restore.py

Deploy: put this file, tape_restore.py, requirements.txt, and README.md
in the root of a Hugging Face Space with SDK = Gradio.
"""

import os
import tempfile
import traceback

import gradio as gr
import numpy as np

import tape_restore as tr


class Args:
    """Lightweight stand-in for the argparse.Namespace tape_restore.restore() expects."""
    pass


def run_restoration(
    audio_path,
    use_noise_sample,
    noise_start,
    noise_end,
    hiss_strength,
    click_threshold,
    hum_freq,
    wow_flutter,
    do_declick,
    do_dehum,
    do_dehiss,
    do_eq,
    do_dynamics,
    progress=gr.Progress(track_tqdm=False),
):
    if audio_path is None:
        raise gr.Error("Upload an audio file first.")

    try:
        progress(0.05, desc="Loading audio...")
        audio, sr = tr.load_audio(audio_path)

        duration = len(audio) / sr
        if duration > 20 * 60:
            raise gr.Error(
                f"File is {duration/60:.1f} minutes long. This demo Space caps "
                f"input at 20 minutes to avoid CPU timeouts — for longer masters, "
                f"run tape_restore.py locally instead."
            )

        args = Args()
        args.no_declick = not do_declick
        args.click_threshold = click_threshold
        args.no_hum = not do_dehum
        args.hum_freq = float(hum_freq)
        args.no_hiss = not do_dehiss
        args.hiss_strength = hiss_strength
        args.noise_sample = (
            [noise_start, noise_end]
            if use_noise_sample and noise_end > noise_start
            else None
        )
        args.wow_flutter = wow_flutter
        args.no_eq = not do_eq
        args.no_dynamics = not do_dynamics

        progress(0.2, desc="Restoring (this can take a while on CPU)...")
        restored = tr.restore(audio, sr, args)

        progress(0.9, desc="Writing output...")
        out_path = os.path.join(
            tempfile.gettempdir(), f"restored_{next(tempfile._get_candidate_names())}.wav"
        )
        tr.save_audio(out_path, restored, sr)

        progress(1.0, desc="Done")
        return out_path, "Restoration complete."

    except gr.Error:
        raise
    except Exception as e:
        traceback.print_exc()
        raise gr.Error(f"Restoration failed: {e}")


with gr.Blocks(title="Tape Restoration") as demo:
    gr.Markdown(
        """
        # Analog Tape Restoration
        Upload audio digitized from a degraded analog master (soundtrack cue,
        film master, etc.) and restore it: de-click, de-hum, de-hiss, EQ
        rebuild, and gentle dynamics restoration.

        **Tip:** for the biggest quality gain, mark a 1–2 second region of
        pure tape hiss / room tone (no music) below — usually found at the
        very head or tail of the reel — so de-hiss learns *this tape's*
        actual noise floor instead of guessing.

        Long files run slowly on free CPU hardware; this demo caps input at
        20 minutes. For longer reels, run `tape_restore.py` locally.
        """
    )

    with gr.Row():
        with gr.Column():
            audio_in = gr.Audio(label="Input audio", type="filepath")

            gr.Markdown("### Noise sample (optional but recommended)")
            use_noise_sample = gr.Checkbox(
                label="Use a noise-only region to profile hiss", value=False
            )
            with gr.Row():
                noise_start = gr.Number(label="Start (seconds)", value=0.0)
                noise_end = gr.Number(label="End (seconds)", value=2.0)

            gr.Markdown("### Stages")
            with gr.Row():
                do_declick = gr.Checkbox(label="De-click / de-pop", value=True)
                do_dehum = gr.Checkbox(label="De-hum", value=True)
            with gr.Row():
                do_dehiss = gr.Checkbox(label="De-hiss", value=True)
                do_eq = gr.Checkbox(label="EQ restoration", value=True)
            do_dynamics = gr.Checkbox(label="Dynamics restoration", value=True)

            gr.Markdown("### Parameters")
            hiss_strength = gr.Slider(
                0.0, 1.0, value=0.75, step=0.05,
                label="De-hiss strength",
                info="Higher = more aggressive noise reduction, but can thin out highs/room tone",
            )
            click_threshold = gr.Slider(
                4.0, 40.0, value=15.0, step=1.0,
                label="Click sensitivity threshold",
                info="Lower = catches more clicks but risks false positives on transients",
            )
            hum_freq = gr.Radio(
                choices=["60", "50"], value="60",
                label="Mains hum frequency",
                info="60 Hz = US/NA tapes, 50 Hz = EU/most of the rest of the world",
            )
            wow_flutter = gr.Slider(
                0.0, 0.5, value=0.0, step=0.05,
                label="Wow/flutter smoothing",
                info="Mild envelope-based warble reduction. 0 = off. For real pitch "
                     "correction, use a dedicated tool like Capstan or iZotope RX.",
            )

            run_btn = gr.Button("Restore", variant="primary")

        with gr.Column():
            audio_out = gr.Audio(label="Restored audio", type="filepath")
            status = gr.Textbox(label="Status", interactive=False)

    run_btn.click(
        fn=run_restoration,
        inputs=[
            audio_in,
            use_noise_sample,
            noise_start,
            noise_end,
            hiss_strength,
            click_threshold,
            hum_freq,
            wow_flutter,
            do_declick,
            do_dehum,
            do_dehiss,
            do_eq,
            do_dynamics,
        ],
        outputs=[audio_out, status],
    )

if __name__ == "__main__":
    demo.launch()