import gradio as gr


_CSS = """
.audio-gallery-container {
    padding: 16px;
}
.audio-gallery-grid {
    display: grid;
    gap: 16px;
}
.audio-item {
    background: var(--block-background-fill, #1e1e2e);
    border: 1px solid var(--block-border-color, #3a3a5c);
    border-radius: 8px;
    padding: 12px;
    display: flex;
    flex-direction: column;
    gap: 8px;
}
.audio-label {
    font-weight: 600;
    font-size: 0.9rem;
    color: var(--body-text-color, #cdd6f4);
    text-transform: uppercase;
    letter-spacing: 0.05em;
}
.waveform-canvas {
    width: 100%;
    height: 60px;
    border-radius: 4px;
    background: var(--background-fill-secondary, #181825);
    display: block;
}
.audio-controls {
    display: flex;
    align-items: center;
    gap: 8px;
}
.audio-action-stack {
    display: flex;
    flex-direction: column;
    align-items: center;
    gap: 6px;
}
.play-btn {
    background: #4a9eff;
    border: none;
    border-radius: 50%;
    width: 32px;
    height: 32px;
    cursor: pointer;
    font-size: 0.85rem;
    color: white;
    flex-shrink: 0;
}
.play-btn:hover {
    background: #6ab4ff;
}
.download-link {
    color: #4a9eff;
    font-size: 0.75rem;
    font-weight: 600;
    text-decoration: none;
}
.download-link:hover {
    color: #6ab4ff;
    text-decoration: underline;
}
.time-display {
    font-size: 0.8rem;
    color: var(--body-text-color, #a6adc8);
    font-family: monospace;
}
"""

GALLERY_JS = """
function formatTime(secs) {
    var m = Math.floor(secs / 60);
    var s = Math.floor(secs % 60).toString().padStart(2, '0');
    return m + ':' + s;
}

function drawWaveform(canvas) {
    var ctx = canvas.getContext('2d');
    var w = canvas.offsetWidth || 300;
    canvas.width = w;
    var h = canvas.height;
    ctx.clearRect(0, 0, w, h);
    ctx.fillStyle = '#4a9eff';
    var bars = 60;
    for (var i = 0; i < bars; i++) {
        var x = (i / bars) * w;
        var bw = Math.max(1, w / bars - 2);
        var amp = h * (0.2 + 0.7 * Math.abs(Math.sin(i * 0.45 + Math.random() * 0.3)));
        var y = (h - amp) / 2;
        ctx.fillRect(x, y, bw, amp);
    }
}

function initAudioItem(item) {
    if (item.getAttribute('data-initialized') === 'true') return;
    item.setAttribute('data-initialized', 'true');
    var audio = item.querySelector('audio');
    var canvas = item.querySelector('.waveform-canvas');
    var btn = item.querySelector('.play-btn');
    var timeDisplay = item.querySelector('.time-display');
    if (!audio || !canvas || !btn) return;

    drawWaveform(canvas);

    btn.addEventListener('click', function () {
        document.querySelectorAll('.audio-item audio').forEach(function (a) {
            if (a !== audio && !a.paused) {
                a.pause();
                a.closest('.audio-item').querySelector('.play-btn').textContent = '\\u25B6';
            }
        });
        if (audio.paused) {
            audio.play();
            btn.textContent = '\\u23F8';
        } else {
            audio.pause();
            btn.textContent = '\\u25B6';
        }
    });

    audio.addEventListener('timeupdate', function () {
        timeDisplay.textContent = formatTime(audio.currentTime);
    });

    audio.addEventListener('ended', function () {
        btn.textContent = '\\u25B6';
    });
}

// Auto-initialize new audio items as they are injected into the DOM
// Watch the entire document body since the gallery doesn't exist on page load
(function setupObserver() {
    document.querySelectorAll('.audio-item').forEach(initAudioItem);

    var observer = new MutationObserver(function (mutations) {
        mutations.forEach(function (m) {
            m.addedNodes.forEach(function (node) {
                if (node.nodeType !== 1) return; // element node only
                if (node.classList && node.classList.contains('audio-item')) {
                    initAudioItem(node);
                }
                if (node.querySelectorAll) {
                    node.querySelectorAll('.audio-item').forEach(initAudioItem);
                }
            });
        });
    });
    observer.observe(document.body, { childList: true, subtree: true });
})();
"""

class AudioGallery(gr.HTML):
    """Gradio HTML component that renders audio stems in a responsive grid."""

    DEFAULT_LABELS = ["Drums", "Vocals", "Guitar", "Bass", "Other", "Piano", "Music", "Full"]

    def __init__(
        self,
        audio_urls,
        *,
        value=None,
        labels=None,
        columns=3,
        label=None,
        **kwargs,
    ):
        labels = labels or self.DEFAULT_LABELS
        html = self._build_html(audio_urls, labels=labels, columns=columns)
        super().__init__(value=html, label=label, **kwargs)

    @staticmethod
    def _build_html(audio_urls, labels, columns):
        items = ""
        for i, url in enumerate(audio_urls):
            lbl = labels[i] if i < len(labels) else f"Track {i + 1}"
            items += (
                f'<div class="audio-item" data-index="{i}" data-initialized="false">'
                f'<div class="audio-label">{lbl}</div>'
                f'<canvas class="waveform-canvas" width="300" height="60"></canvas>'
                f'<audio src="{url}" preload="metadata"></audio>'
                f'<div class="audio-controls">'
                f'<div class="audio-action-stack">'
                f'<button class="play-btn" type="button">&#9654;</button>'
                f'<a class="download-link" href="{url}" download>Download</a>'
                f'</div>'
                f'<div class="time-display">0:00</div>'
                f'</div>'
                f'</div>\n'
            )
        return (
            f'<style>{_CSS}</style>'
            f'<div class="audio-gallery-container">'
            f'<div class="audio-gallery-grid" style="grid-template-columns: repeat({columns}, 1fr);">'
            f'{items}'
            f'</div>'
            f'</div>'
        )