// static/js/ui/tts.js

import { dom } from './dom.js';

const ttsCache = new Map();

const ttsState = {
    currentMessageIndex: null,
    animationFrameId: null,
    activeButton: null,
};

const ttsStreamManager = {
    HF_WEBSOCKET_URL: "wss://ezmary-ttslive.hf.space/ws",
    socket: null,
    activeStreamController: null,
    pendingRequest: null, // برای نگهداری درخواست در زمان اتصال مجدد

    connectWebSocket() {
        // فقط اگر اتصالی وجود ندارد یا قطع شده، اتصال جدید بساز
        if (this.socket && this.socket.readyState !== WebSocket.CLOSED) {
            return;
        }

        this.socket = new WebSocket(this.HF_WEBSOCKET_URL);

        this.socket.onopen = () => {
            console.log("اتصال WebSocket به سرور TTS برقرار شد.");
            // اگر درخواستی در صف انتظار بود، آن را ارسال کن
            if (this.pendingRequest) {
                const { messageIndex, text, button } = this.pendingRequest;
                this.pendingRequest = null;
                this.stream(messageIndex, text, button, true); // ارسال مجدد درخواست
            }
        };
        
        this.socket.onmessage = (event) => {
            if (this.activeStreamController && this.activeStreamController.handleMessage) {
                this.activeStreamController.handleMessage(event);
            }
        };

        this.socket.onclose = () => {
            console.warn("اتصال WebSocket قطع شد.");
            this.socket = null; // سوکت را null کن تا اتصال بعدی دوباره برقرار شود
            if (this.activeStreamController) {
                this.activeStreamController.stop();
            }
        };

        this.socket.onerror = (error) => {
            console.error("خطای WebSocket:", error);
            if (this.socket) this.socket.close();
        };
    },
    
    _createStreamController(messageIndex, button) {
        // ... (این تابع داخلی بدون تغییر باقی می‌ماند)
        const controller = {
            messageIndex: messageIndex, button: button, audioContext: null, audioQueue: [], receivedPcmChunks: [], sourceNodes: [], isStopped: false, isPlaying: false, nextStartTime: 0, timerInterval: null, startTime: 0, elapsedTime: 0,
            _initializeAudio() { if (!this.audioContext || this.audioContext.state === 'closed') { this.audioContext = new (window.AudioContext || window.webkitAudioContext)({ sampleRate: 24000 }); } if(this.audioContext.state === 'suspended') { this.audioContext.resume(); } this.nextStartTime = this.audioContext.currentTime; },
            async handleMessage(event) { if (this.isStopped) return; if (typeof event.data === 'string') { const message = JSON.parse(event.data); if (message.event === "STREAM_ENDED") this._handleStreamEnd(); else if (message.event === "ERROR") { console.error(`خطا از سرور TTS: ${message.message}`); this.stop(); } } else { const arrayBuffer = await event.data.arrayBuffer(); const pcmData = new Int16Array(arrayBuffer); this.audioQueue.push(pcmData); this.receivedPcmChunks.push(pcmData); if (!this.isPlaying) this._playFromQueue(); } },
            async _playFromQueue() { if (this.audioQueue.length === 0 || this.isStopped) { this.isPlaying = false; return; } this.isPlaying = true; if (!dom.globalAudioPlayer.classList.contains('visible')) { showGlobalPlayer(false, true); this.startTime = Date.now(); this.timerInterval = setInterval(() => updateGlobalPlayerUI(this), 250); ttsState.animationFrameId = requestAnimationFrame(drawWaveform); if(this.button) { this.button.classList.remove('loading'); this.button.classList.add('playing'); } } while (this.audioQueue.length > 0) { if (this.isStopped) break; const pcmData = this.audioQueue.shift(); const float32Data = new Float32Array(pcmData.length); for (let i = 0; i < pcmData.length; i++) float32Data[i] = pcmData[i] / 32768.0; if (this.audioContext.state === 'closed') return; const audioBuffer = this.audioContext.createBuffer(1, float32Data.length, this.audioContext.sampleRate); audioBuffer.getChannelData(0).set(float32Data); const source = this.audioContext.createBufferSource(); source.buffer = audioBuffer; source.connect(this.audioContext.destination); const currentTime = this.audioContext.currentTime; this.nextStartTime = Math.max(this.nextStartTime, currentTime); source.start(this.nextStartTime); this.sourceNodes.push(source); this.nextStartTime += audioBuffer.duration; } this.isPlaying = false; },
            _handleStreamEnd() { this._finalizeAndCacheAudio(); const checkPlaybackEnd = setInterval(() => { if (this.audioQueue.length === 0 && this.audioContext && this.audioContext.currentTime > this.nextStartTime - 0.1) { if(!this.isStopped) this.stop(); clearInterval(checkPlaybackEnd); } }, 100); },
            _finalizeAndCacheAudio() { if (this.receivedPcmChunks.length === 0) return; const totalLength = this.receivedPcmChunks.reduce((acc, val) => acc + val.length, 0); const concatenatedPcm = new Int16Array(totalLength); let offset = 0; for (const chunk of this.receivedPcmChunks) { concatenatedPcm.set(chunk, offset); offset += chunk.length; } const wavBlob = this._createWavBlob(concatenatedPcm, this.audioContext.sampleRate); ttsCache.set(this.messageIndex, wavBlob); this.receivedPcmChunks = []; },
            _createWavBlob(pcmData, sampleRate) { const numChannels = 1, bitsPerSample = 16; const blockAlign = (numChannels * bitsPerSample) / 8; const byteRate = sampleRate * blockAlign; const dataSize = pcmData.length * (bitsPerSample / 8); const buffer = new ArrayBuffer(44 + dataSize); const view = new DataView(buffer); function writeString(view, offset, string) { for (let i = 0; i < string.length; i++) view.setUint8(offset + i, string.charCodeAt(i)); } writeString(view, 0, 'RIFF'); view.setUint32(4, 36 + dataSize, true); writeString(view, 8, 'WAVE'); writeString(view, 12, 'fmt '); view.setUint32(16, 16, true); view.setUint16(20, 1, true); view.setUint16(22, numChannels, true); view.setUint32(24, sampleRate, true); view.setUint32(28, byteRate, true); view.setUint16(32, blockAlign, true); view.setUint16(34, bitsPerSample, true); writeString(view, 36, 'data'); view.setUint32(40, dataSize, true); for (let i = 0; i < pcmData.length; i++) view.setInt16(44 + i * 2, pcmData[i], true); return new Blob([view], { type: 'audio/wav' }); },
            stop() { this.isStopped = true; if (this.timerInterval) clearInterval(this.timerInterval); this.timerInterval = null; this.sourceNodes.forEach(source => { try { source.stop(); } catch(e) {} }); if (this.audioContext && this.audioContext.state !== 'closed') { this.audioContext.close(); } if (ttsStreamManager.activeStreamController === this) { ttsStreamManager.activeStreamController = null; hideGlobalPlayer(); if (this.button) this.button.classList.remove('playing', 'loading'); ttsState.activeButton = null; } },
            pause() { if (!this.audioContext || this.audioContext.state !== 'running') return; this.audioContext.suspend().then(() => { clearInterval(this.timerInterval); this.timerInterval = null; this.elapsedTime += (Date.now() - this.startTime) / 1000; showGlobalPlayer(false, false); if(this.button) this.button.classList.remove('playing'); }); },
            resume() { if (!this.audioContext || this.audioContext.state !== 'suspended') return; this.audioContext.resume().then(() => { this.startTime = Date.now(); this.timerInterval = setInterval(() => updateGlobalPlayerUI(this), 250); showGlobalPlayer(false, true); if(this.button) this.button.classList.add('playing'); }); }
        };
        controller._initializeAudio();
        return controller;
    },

    stopCurrentStream() {
        if (this.activeStreamController) {
            this.activeStreamController.stop();
            this.activeStreamController = null;
        }
        // *** مهم: اتصال وب‌سوکت را قطع کن ***
        if (this.socket && this.socket.readyState === WebSocket.OPEN) {
            this.socket.close();
            console.log("اتصال WebSocket به دلیل توقف توسط کاربر، عمداً قطع شد.");
        }
        this.pendingRequest = null; // هر درخواست در حال انتظاری را لغو کن
    },

    stream(messageIndex, text, button, isRetrying = false) {
        if (!isRetrying) {
            this.stopCurrentStream(); // همیشه پخش قبلی را کاملا متوقف کن
        }
        
        // اگر اتصال برقرار نیست، برقرار کن و درخواست را در صف بگذار
        if (!this.socket || this.socket.readyState !== WebSocket.OPEN) {
            this.pendingRequest = { messageIndex, text, button };
            setSpeakButtonLoading(button, true);
            this.connectWebSocket(); // این تابع اتصال جدید را شروع می‌کند
            return;
        }

        this.activeStreamController = this._createStreamController(messageIndex, button);
        ttsState.currentMessageIndex = messageIndex;
        ttsState.activeButton = button;
        
        this.socket.send(text);
        setSpeakButtonLoading(button, true);
    },
    
    async playFromCache(messageIndex, button) {
        const audioBlob = ttsCache.get(messageIndex);
        if (!audioBlob) return;
        
        this.stopCurrentStream();
        this.activeStreamController = this._createStreamController(messageIndex, button);
        ttsState.currentMessageIndex = messageIndex;
        ttsState.activeButton = button;
        
        try {
            const controller = this.activeStreamController;
            const arrayBuffer = await audioBlob.arrayBuffer();
            if (controller.audioContext.state === 'closed') return;
            const decodedBuffer = await controller.audioContext.decodeAudioData(arrayBuffer);
            const source = controller.audioContext.createBufferSource();
            source.buffer = decodedBuffer;
            source.connect(controller.audioContext.destination);
            
            showGlobalPlayer(false, true);
            button.classList.add('playing');
            
            controller.startTime = Date.now();
            controller.timerInterval = setInterval(() => updateGlobalPlayerUI(controller), 250);
            ttsState.animationFrameId = requestAnimationFrame(drawWaveform);

            source.start(0);
            source.onended = () => {
                if (controller.audioContext && controller.audioContext.state !== 'closed') {
                   controller.stop();
                }
            };
            controller.sourceNodes.push(source);
        } catch (error) {
            console.error("خطا در پخش از کش:", error);
            this.stopCurrentStream();
        }
    }
};

export function clearCacheForMessage(messageIndex) { ttsCache.delete(messageIndex); }
export function clearAllCache() { ttsCache.clear(); }
export function hasCacheForMessage(messageIndex) { return ttsCache.has(messageIndex); }
export function getAudioState() { const controller = ttsStreamManager.activeStreamController; if (!controller || !controller.audioContext) return { status: 'idle', messageIndex: null }; return { status: controller.audioContext.state, messageIndex: controller.messageIndex }; }
export function togglePauseResumeAudio() { const controller = ttsStreamManager.activeStreamController; if (!controller) return; if (controller.audioContext.state === 'running') controller.pause(); else if (controller.audioContext.state === 'suspended') controller.resume(); }
export function stopAudio() { ttsStreamManager.stopCurrentStream(); }
export function stream(messageIndex, text, button) { ttsStreamManager.stream(messageIndex, text, button); }
export function playFromCache(messageIndex, button) { ttsStreamManager.playFromCache(messageIndex, button); }
export function initTtsPlayer() { dom.globalPlayerPlayPause.addEventListener('click', togglePauseResumeAudio); dom.globalPlayerClose.addEventListener('click', stopAudio); }
function formatTime(s) { if (isNaN(s) || s < 0) return '0:00'; const minutes = Math.floor(s / 60); const seconds = Math.floor(s % 60).toString().padStart(2, '0'); return `${minutes}:${seconds}`; }
function drawWaveform() { if (!dom.waveformCanvas || !ttsStreamManager.activeStreamController) { if (ttsState.animationFrameId) { cancelAnimationFrame(ttsState.animationFrameId); ttsState.animationFrameId = null; } return; } const canvas = dom.waveformCanvas; const ctx = canvas.getContext('2d'); const dpr = window.devicePixelRatio || 1; const rect = canvas.getBoundingClientRect(); if (canvas.width !== rect.width * dpr || canvas.height !== rect.height * dpr) { canvas.width = rect.width * dpr; canvas.height = rect.height * dpr; ctx.scale(dpr, dpr); } const width = canvas.width / dpr, height = canvas.height / dpr; ctx.clearRect(0, 0, width, height); const barWidth = 2, barGap = 1.5, totalBarWidth = barWidth + barGap; const numBars = Math.floor(width / totalBarWidth); const offset = (width - numBars * totalBarWidth) / 2; const inactiveColor = getComputedStyle(document.documentElement).classList.contains('dark') ? '#4b5563' : '#d1d5db'; const controller = ttsStreamManager.activeStreamController; for (let i = 0; i < numBars; i++) { let barHeight = height * 0.1; if(controller && controller.audioContext && controller.audioContext.state === 'running') { barHeight = (Math.sin((i + Date.now() / 200) * 0.2) + 1) / 2 * height * 0.7 + height * 0.1; } const x = offset + i * totalBarWidth, y = (height - barHeight) / 2; ctx.fillStyle = inactiveColor; ctx.fillRect(x, y, barWidth, barHeight); } ttsState.animationFrameId = requestAnimationFrame(drawWaveform); }
function updateGlobalPlayerUI(controller) { if (!controller || !controller.audioContext || controller.audioContext.state !== 'running') return; const currentTime = controller.elapsedTime + (Date.now() - controller.startTime) / 1000; dom.globalPlayerCurrentTime.textContent = formatTime(currentTime); }
function showGlobalPlayer(isLoading = false, isPlaying = false) { dom.globalAudioPlayer.classList.add('visible'); dom.globalPlayerText.classList.add('hidden'); dom.globalPlayerTotalTime.classList.add('hidden'); dom.globalPlayerPlayPause.style.display = isLoading ? 'none' : 'flex'; dom.globalPlayerLoading.classList.toggle('hidden', !isLoading); dom.globalPlayerPlayIcon.classList.toggle('hidden', isPlaying); dom.globalPlayerPauseIcon.classList.toggle('hidden', !isPlaying); }
export function hideGlobalPlayer() { dom.globalAudioPlayer.classList.remove('visible'); ttsState.currentMessageIndex = null; if (ttsState.animationFrameId) { cancelAnimationFrame(ttsState.animationFrameId); ttsState.animationFrameId = null; } }
export function setSpeakButtonLoading(button, isLoading) { if (button) { button.classList.toggle('loading', isLoading); if(isLoading) button.classList.remove('playing'); } }