const express = require('express');
const WebSocket = require('ws');
const { spawn } = require('child_process');

const app = express();
const PORT = 3000;

app.use(express.static('public'));

const server = app.listen(PORT, () => {
    console.log(`Server running on http://localhost:${PORT}`);
});

const wss = new WebSocket.Server({ server });

class SpeechProcessor {
    constructor() {
        this.pythonProcess = null;
        this.activeRequests = new Map();
        this.requestCounter = 0;
        this.initializePythonProcess();
        this.buffer = Buffer.alloc(0);
    }

    initializePythonProcess() {
        this.pythonProcess = spawn('python3', ['speech_processor.py']);

        this.pythonProcess.stderr.on('data', (data) => {
            console.error('Python STDERR:', data.toString());
        });

        this.pythonProcess.on('close', (code) => {
            console.log(`Python process exited with code ${code}`);
            this.activeRequests.clear();
            setTimeout(() => this.initializePythonProcess(), 1000);
        });

        this.pythonProcess.stdout.on('data', (data) => {
            this.buffer = Buffer.concat([this.buffer, data]);
            this.processBuffer();
        });
    }

    processBuffer() {
        while (this.buffer.length >= 8) {
            const length = this.buffer.readUInt32BE(0);
            const requestId = this.buffer.readUInt32BE(4);

            if (this.buffer.length >= 8 + length) {
                const message = this.buffer.slice(8, 8 + length);
                this.buffer = this.buffer.slice(8 + length);

                try {
                    const result = JSON.parse(message.toString());
                    if (this.activeRequests.has(requestId)) {
                        const { resolve } = this.activeRequests.get(requestId);
                        this.activeRequests.delete(requestId);
                        resolve(result);
                    }
                } catch (error) {
                    console.error('Failed to parse message:', error);
                }
            } else {
                break;
            }
        }
    }

    async processAudio(audioBuffer) {
        return new Promise((resolve, reject) => {
            if (!this.pythonProcess) {
                reject(new Error('Processor not ready'));
                return;
            }

            const requestId = this.requestCounter++;
            this.activeRequests.set(requestId, { resolve, reject });

            const lengthBuffer = Buffer.alloc(4);
            lengthBuffer.writeUInt32BE(audioBuffer.length, 0);
            
            const idBuffer = Buffer.alloc(4);
            idBuffer.writeUInt32BE(requestId, 0);

            this.pythonProcess.stdin.write(lengthBuffer);
            this.pythonProcess.stdin.write(idBuffer);
            this.pythonProcess.stdin.write(audioBuffer);

            setTimeout(() => {
                if (this.activeRequests.has(requestId)) {
                    this.activeRequests.delete(requestId);
                    reject(new Error('Processing timeout'));
                }
            }, 5000);
        });
    }
}

const speechProcessor = new SpeechProcessor();

wss.on('connection', (ws) => {
    console.log('Client connected-C3');
    
    let lastFinalText = '';
    let lastPartialText = '';
    let partialTextTimeout = null;

    ws.on('message', async (message) => {
        if (!Buffer.isBuffer(message)) return;

        try {
            const result = await speechProcessor.processAudio(message);
            
            if (!result.success) {
                console.error('Processing error:', result.error);
                return;
            }

            // Handle final results
            if (result.is_final && result.text && result.text !== lastFinalText) {
                lastFinalText = result.text;
                ws.send(JSON.stringify({
                    type: 'transcription',
                    text: result.text,
                    is_final: true
                }));
                console.log('Final:', result.text);
                lastPartialText = '';
                if (partialTextTimeout) {
                    clearTimeout(partialTextTimeout);
                    partialTextTimeout = null;
                }
            } 
            // Handle partial results with debouncing
            else if (!result.is_final && result.text && result.text !== lastPartialText) {
                lastPartialText = result.text;
                
                if (partialTextTimeout) {
                    clearTimeout(partialTextTimeout);
                }
                
                partialTextTimeout = setTimeout(() => {
                    ws.send(JSON.stringify({
                        type: 'partial_transcription',
                        text: result.text,
                        is_final: false
                    }));
                    partialTextTimeout = null;
                }, 300); // 300ms debounce
            }
        } catch (error) {
            console.error('Processing error:', error);
        }
    });

    ws.on('close', () => {
        console.log('Client disconnected');
        if (partialTextTimeout) {
            clearTimeout(partialTextTimeout);
        }
    });
});