const ws = require('ws');

class ColorifyAI {
    _createSocket(endpoint, data) {
        return new Promise((resolve, reject) => {
            const sessionHash = Math.random().toString(36).substring(2);
            const socket = new ws(`wss://colorifyai.art/${endpoint}/queue/join`);
            
            socket.on('message', (msg) => {
                const d = JSON.parse(msg.toString('utf8'));
                
                switch (d.msg) {
                    case 'send_hash':
                        socket.send(JSON.stringify({ session_hash: sessionHash }));
                        break;
                    case 'send_data':
                        socket.send(JSON.stringify({ data }));
                        break;
                    case 'process_completed':
                        socket.close();
                        resolve(d.output.result[0]);
                        break;
                    case 'estimation':
                    case 'process_starts':
                        break;
                    default:
                        console.log(`Unexpected message: ${msg.toString('utf8')}`);
                }
            });
            
            socket.on('error', reject);
        });
    }

    text2sketch = async (prompt, { ratio = '1:1', style = 'default' } = {}) => {
        const _ratio = ['1:1', '3:4', '4:3', '9:16', '16:9', '2:3', '3:2'];
        const _style = ['default', 'sci-fi', 'pixel', 'chibi', 'graffiti', 'minimalist', 'anime'];
        
        if (!prompt) throw new Error('Prompt is required');
        if (!_ratio.includes(ratio)) throw new Error(`Available ratios: ${_ratio.join(', ')}`);
        if (!_style.includes(style)) throw new Error(`Available styles: ${_style.join(', ')}`);
        
        return this._createSocket('demo-colorify-text2img', {
            aspect_ratio: ratio,
            prompt,
            request_from: 10,
            style
        });
    }
    
    image2sketch = async (buffer) => {
        if (!buffer || !Buffer.isBuffer(buffer)) throw new Error('Image buffer is required');
        
        const result = await this._createSocket('demo-colorify-img2img', {
            request_from: 10,
            source_image: `data:image/jpeg;base64,${buffer.toString('base64')}`
        });
        return `https://temp.colorifyai.art/${result}`;
    }
}

// Usage:
const c = new ColorifyAI();
const resp = await c.text2sketch('a pikachu', { ratio: '16:9', style: 'anime' });
console.log(resp);