import gradio as gr
import os
import uuid
from pydub import AudioSegment, effects
from pydub.silence import split_on_silence
import re
import time
import numpy as np


# ─────────────────────────────────────────────
#  UTILS
# ─────────────────────────────────────────────

def clean_file_name(file_path):
    file_name = os.path.basename(file_path)
    file_name, file_extension = os.path.splitext(file_name)
    cleaned = re.sub(r'[^a-zA-Z\d]+', '_', file_name)
    clean_name = re.sub(r'_+', '_', cleaned).strip('_')
    random_uuid = uuid.uuid4().hex[:6]
    clean_file_path = os.path.join(
        os.path.dirname(file_path),
        clean_name + f"_{random_uuid}" + file_extension
    )
    return clean_file_path


def calculate_duration(file_path):
    audio = AudioSegment.from_file(file_path)
    return len(audio) / 1000.0


# ─────────────────────────────────────────────
#  FILE TRACKING / CLEANUP
# ─────────────────────────────────────────────

FILE_TIMESTAMPS = {}

def track_file(file_path):
    FILE_TIMESTAMPS[file_path] = time.time()


def cleanup_tracked_files(max_age_seconds=3600):
    now = time.time()
    to_delete = []
    for file_path, created_time in FILE_TIMESTAMPS.items():
        if now - created_time > max_age_seconds:
            if os.path.exists(file_path):
                try:
                    os.remove(file_path)
                    print(f"Deleted: {file_path}")
                except Exception as e:
                    print(f"Error deleting {file_path}: {e}")
            to_delete.append(file_path)
    for f in to_delete:
        FILE_TIMESTAMPS.pop(f, None)


def start_cleanup_worker(interval=3600):
    def worker():
        while True:
            print("Cleaning tracked files...")
            cleanup_tracked_files()
            time.sleep(interval)
    import threading
    threading.Thread(target=worker, daemon=True).start()


# ─────────────────────────────────────────────
#  CORE FUNCTIONS
# ─────────────────────────────────────────────

def remove_silence(file_path, minimum_silence=50):
    sound = AudioSegment.from_file(file_path)
    audio_chunks = split_on_silence(
        sound,
        min_silence_len=100,
        silence_thresh=-45,
        keep_silence=minimum_silence
    )
    combined = AudioSegment.empty()
    for chunk in audio_chunks:
        combined += chunk
    output_path = clean_file_name(file_path)
    combined.export(output_path)
    return output_path


def process_audio(audio_file, seconds=0.05):
    if audio_file is None:
        return None, None, "No file uploaded", None
    if not os.path.exists(audio_file):
        return None, None, "File not found", None

    track_file(audio_file)
    keep_silence = int(seconds * 1000)
    output_audio_file = remove_silence(audio_file, minimum_silence=keep_silence)
    track_file(output_audio_file)

    before = calculate_duration(audio_file)
    after  = calculate_duration(output_audio_file)
    text   = f"Old Duration: {before:.3f} seconds\nNew Duration: {after:.3f} seconds"

    # 4th value stored in State so enhancer can access it
    return output_audio_file, output_audio_file, text, output_audio_file


def enhance_audio(file_path):
    """
    Enhancement Pipeline:
      1. Noise Reduction  (noisereduce, if installed — optional)
      2. High-Pass Filter (scipy, removes rumble below 80 Hz)
      3. Normalization    (pydub effects.normalize)
    Gracefully skips unavailable libraries.
    """
    if file_path is None:
        return None, None, "No audio to enhance. Please run Remove Silence first."
    if not os.path.exists(file_path):
        return None, None, "File not found."

    sound        = AudioSegment.from_file(file_path)
    sample_rate  = sound.frame_rate
    channels     = sound.channels
    sample_width = sound.sample_width

    # Convert to float32 in range [-1, 1]
    raw = np.array(sound.get_array_of_samples(), dtype=np.float32)
    if sample_width == 1:
        raw = (raw - 128.0) / 128.0
    elif sample_width == 2:
        raw = raw / 32768.0
    elif sample_width == 4:
        raw = raw / 2147483648.0

    if channels == 2:
        raw = raw.reshape((-1, 2))

    # Step 1 — Noise Reduction
    try:
        import noisereduce as nr
        if channels == 2:
            raw[:, 0] = nr.reduce_noise(y=raw[:, 0], sr=sample_rate, prop_decrease=0.75)
            raw[:, 1] = nr.reduce_noise(y=raw[:, 1], sr=sample_rate, prop_decrease=0.75)
        else:
            raw = nr.reduce_noise(y=raw, sr=sample_rate, prop_decrease=0.75)
    except Exception:
        pass

    # Step 2 — High-Pass Filter (remove rumble below 80 Hz)
    try:
        from scipy.signal import butter, sosfilt
        sos = butter(4, 80, btype='hp', fs=sample_rate, output='sos')
        if channels == 2:
            raw[:, 0] = sosfilt(sos, raw[:, 0])
            raw[:, 1] = sosfilt(sos, raw[:, 1])
        else:
            raw = sosfilt(sos, raw)
    except Exception:
        pass

    if channels == 2:
        raw = raw.flatten()

    raw         = np.clip(raw, -1.0, 1.0)
    out_samples = (raw * 32768.0).astype(np.int16)

    enhanced_segment = AudioSegment(
        out_samples.tobytes(),
        frame_rate   = sample_rate,
        sample_width = 2,
        channels     = channels
    )

    # Step 3 — Normalize volume
    enhanced_segment = effects.normalize(enhanced_segment)

    base        = os.path.splitext(clean_file_name(file_path))[0]
    output_path = base + "_enhanced.wav"
    enhanced_segment.export(output_path, format="wav")
    track_file(output_path)

    duration = calculate_duration(output_path)
    text     = f"Enhanced Successfully!\nDuration: {duration:.3f} seconds"
    return output_path, output_path, text


# ─────────────────────────────────────────────
#  UI
# ─────────────────────────────────────────────

def ui():
    theme = gr.themes.Soft(
        font=[gr.themes.GoogleFont("Source Sans Pro"), "Arial", "sans-serif"]
    )

    css = """
    .gradio-container { max-width: none !important; }
    .tab-content { padding: 20px; }

    @keyframes rgb-glow {
        0%   { color: #ff0080; text-shadow: 0 0 12px #ff0080, 0 0 24px #ff0080; }
        16%  { color: #ff8c00; text-shadow: 0 0 12px #ff8c00, 0 0 24px #ff8c00; }
        33%  { color: #ffe600; text-shadow: 0 0 12px #ffe600, 0 0 24px #ffe600; }
        50%  { color: #00ff88; text-shadow: 0 0 12px #00ff88, 0 0 24px #00ff88; }
        66%  { color: #00c8ff; text-shadow: 0 0 12px #00c8ff, 0 0 24px #00c8ff; }
        83%  { color: #a855f7; text-shadow: 0 0 12px #a855f7, 0 0 24px #a855f7; }
        100% { color: #ff0080; text-shadow: 0 0 12px #ff0080, 0 0 24px #ff0080; }
    }
    @keyframes rgb-border {
        0%   { border-color: #ff0080; box-shadow: 0 0 10px #ff0080; }
        16%  { border-color: #ff8c00; box-shadow: 0 0 10px #ff8c00; }
        33%  { border-color: #ffe600; box-shadow: 0 0 10px #ffe600; }
        50%  { border-color: #00ff88; box-shadow: 0 0 10px #00ff88; }
        66%  { border-color: #00c8ff; box-shadow: 0 0 10px #00c8ff; }
        83%  { border-color: #a855f7; box-shadow: 0 0 10px #a855f7; }
        100% { border-color: #ff0080; box-shadow: 0 0 10px #ff0080; }
    }
    @keyframes gradient-shift {
        0%   { background-position: 0% 50%; }
        50%  { background-position: 100% 50%; }
        100% { background-position: 0% 50%; }
    }
    @keyframes pulse-glow {
        0%, 100% { box-shadow: 0 0 15px rgba(0,200,255,0.45), 0 4px 20px rgba(0,0,0,0.3); }
        50%       { box-shadow: 0 0 38px rgba(168,85,247,0.65), 0 4px 20px rgba(0,0,0,0.3); }
    }

    .rgb-title {
        animation: rgb-glow 3s linear infinite;
        font-weight: 900;
        letter-spacing: 1px;
    }

    .made-by-badge {
        display: inline-block;
        padding: 6px 22px;
        border: 2px solid #00c8ff;
        border-radius: 50px;
        font-size: 0.82em;
        font-weight: 800;
        letter-spacing: 3px;
        text-transform: uppercase;
        animation: rgb-glow 3s linear infinite, rgb-border 3s linear infinite;
        background: rgba(0,0,0,0.04);
        margin-top: 10px;
    }

    /* Remove Silence primary button */
    button.primary {
        background: linear-gradient(135deg, #2563eb, #1e40af) !important;
        color: white !important;
        font-weight: 700 !important;
        border: none !important;
        border-radius: 12px !important;
        padding: 13px 20px !important;
        font-size: 1.05em !important;
        box-shadow: 0 4px 15px rgba(37,99,235,0.4) !important;
        transition: all 0.3s ease !important;
    }
    button.primary:hover {
        background: linear-gradient(135deg, #1d4ed8, #1e3a8a) !important;
        box-shadow: 0 6px 25px rgba(37,99,235,0.65) !important;
        transform: translateY(-1px) !important;
    }

    /* Animated divider */
    .enhance-divider {
        margin: 36px auto 22px;
        max-width: 95%;
        height: 2px;
        background: linear-gradient(90deg, transparent, #00c8ff, #a855f7, #ff0080, transparent);
        background-size: 200% 200%;
        animation: gradient-shift 4s ease infinite;
        border-radius: 2px;
    }

    /* Enhance header card */
    .enhance-header { text-align: center; margin-bottom: 20px; }
    .enhance-label {
        display: inline-flex;
        flex-direction: column;
        align-items: center;
        gap: 4px;
        background: linear-gradient(135deg, #0f172a, #1e293b);
        border: 2px solid #00c8ff;
        border-radius: 16px;
        padding: 12px 32px;
        animation: rgb-border 3s linear infinite;
    }
    .etitle {
        font-size: 1.2em;
        font-weight: 800;
        letter-spacing: 2px;
        text-transform: uppercase;
        animation: rgb-glow 3s linear infinite;
    }
    .esub {
        font-size: 0.75em;
        color: #94a3b8;
        letter-spacing: 1px;
    }

    /* Enhance button */
    .enhance-btn button {
        background: linear-gradient(135deg, #0f172a, #1e293b, #0f172a) !important;
        background-size: 200% 200% !important;
        animation: gradient-shift 4s ease infinite, pulse-glow 2.5s ease-in-out infinite !important;
        color: #00c8ff !important;
        font-weight: 800 !important;
        border: 2px solid #00c8ff !important;
        border-radius: 12px !important;
        padding: 13px 20px !important;
        font-size: 1.05em !important;
        letter-spacing: 1px !important;
        text-transform: uppercase !important;
        transition: color 0.3s ease, background 0.3s ease, transform 0.2s ease !important;
    }
    .enhance-btn button:hover {
        color: #ffffff !important;
        background: linear-gradient(135deg, #00c8ff, #a855f7) !important;
        border-color: transparent !important;
        transform: translateY(-2px) !important;
        box-shadow: 0 8px 30px rgba(0,200,255,0.55) !important;
        animation: none !important;
    }
    """

    with gr.Blocks(theme=theme, css=css) as demo:

        processed_file_state = gr.State(value=None)

        # ── HEADER ──
        gr.HTML("""
        <div style="text-align:center; margin:24px auto 20px; max-width:820px;">
            <h1 class="rgb-title" style="font-size:2.5em; margin-bottom:6px;">
                🔇 Remove Silence From Audio
            </h1>
            <p style="font-size:1.05em; color:#555; margin:0 0 8px;">
                Upload an MP3 or WAV file, and it will remove silent parts from it.
            </p>
            <p style="font-size:0.8em; color:#999; margin-bottom:16px;">
                ⚠️ Please don't upload copyrighted content — it can take this Space offline.
            </p>
            <div class="made-by-badge">✦ Made by Deepu ✦</div>
        </div>
        """)

        # ── REMOVE SILENCE ──
        with gr.Row():
            with gr.Column(scale=1):
                audio_input = gr.Audio(
                    label="📂 Upload Audio",
                    type="filepath",
                    sources=["upload", "microphone"]
                )
                silence_threshold = gr.Number(
                    label="Keep Silence Upto (In seconds)",
                    value=0.05
                )
                submit_btn = gr.Button("🔇 Remove Silence", variant="primary")

            with gr.Column(scale=1):
                audio_output    = gr.Audio(label="▶ Play Audio")
                file_output     = gr.File(label="⬇ Download Audio File")
                duration_output = gr.Textbox(label="📊 Duration Info")

        submit_btn.click(
            fn=process_audio,
            inputs=[audio_input, silence_threshold],
            outputs=[audio_output, file_output, duration_output, processed_file_state]
        )

        # ── DIVIDER ──
        gr.HTML('<div class="enhance-divider"></div>')

        # ── ENHANCE HEADER ──
        gr.HTML("""
        <div class="enhance-header">
            <div class="enhance-label">
                <span class="etitle">✨ Audio Enhancer</span>
                <span class="esub">Noise Reduction &nbsp;·&nbsp; High-Pass Filter &nbsp;·&nbsp; Volume Normalize</span>
            </div>
        </div>
        """)

        # ── ENHANCE SECTION ──
        with gr.Row():
            with gr.Column(scale=1):
                enhance_btn = gr.Button(
                    "✨ Enhance Audio",
                    elem_classes=["enhance-btn"]
                )
                enhance_status = gr.Textbox(
                    label="⚡ Status",
                    interactive=False,
                    placeholder="Run Remove Silence first, then click Enhance..."
                )

            with gr.Column(scale=1):
                enhanced_audio_output = gr.Audio(label="▶ Enhanced Audio")
                enhanced_file_output  = gr.File(label="⬇ Download Enhanced Audio")

        enhance_btn.click(
            fn=enhance_audio,
            inputs=[processed_file_state],
            outputs=[enhanced_audio_output, enhanced_file_output, enhance_status]
        )

    return demo


# ─────────────────────────────────────────────
#  ENTRY POINT
# ─────────────────────────────────────────────

import click

@click.command()
@click.option("--debug", is_flag=True, default=False, help="Enable debug mode.")
@click.option("--share", is_flag=True, default=False, help="Enable sharing of the interface.")
def main(debug, share):
    start_cleanup_worker()
    demo = ui()
    demo.queue().launch(debug=debug, share=share)

if __name__ == "__main__":
    main()