const express = require('express');
const axios = require('axios');

class DeepInfraHandler {
    constructor() {
        this.API_URL = 'https://api.deepinfra.com/v1/openai/chat/completions';
        this.headers = {
            'Accept': 'text/event-stream',
            'Accept-Encoding': 'gzip, deflate, br, zstd',
            'Content-Type': 'application/json',
            'Connection': 'keep-alive',
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36'
        };
    }

    _preparePayload(params) {
        return {
            model: params.model,
            messages: params.messages,
            temperature: params.temperature || 0.7,
            max_tokens: params.max_tokens || 4096,
            top_p: params.top_p || 1.0,
            frequency_penalty: params.frequency_penalty || 0.0,
            presence_penalty: params.presence_penalty || 0.0,
            stop: params.stop || [],
            stream: params.stream || false
        };
    }

    async generateCompletion(params) {
        const payload = this._preparePayload(params);
        
        if (payload.stream) {
            const response = await axios({
                method: 'post',
                url: this.API_URL,
                data: payload,
                headers: this.headers,
                responseType: 'stream'
            });
            return response.data;
        } else {
            const response = await axios.post(this.API_URL, payload, {
                headers: this.headers
            });
            return response.data;
        }
    }
}

const app = express();
const PORT = process.env.PORT || 7860;
const apiHandler = new DeepInfraHandler();

app.use(express.json());

app.post('/chat/completions', async (req, res) => {
    try {
        if (req.body.stream) {
            res.setHeader('Content-Type', 'text/event-stream');
            res.setHeader('Cache-Control', 'no-cache');
            res.setHeader('Connection', 'keep-alive');

            const stream = await apiHandler.generateCompletion(req.body);

            stream.on('data', chunk => {
                const lines = chunk.toString().split('\n');
                for (const line of lines) {
                    if (line.startsWith('data: ')) {
                        try {
                            const content = JSON.parse(line.slice(6));
                            if (content === '[DONE]') {
                                res.write('data: [DONE]\n\n');
                                continue;
                            }
                            res.write(line + '\n\n');
                        } catch (error) {
                            continue;
                        }
                    }
                }
            });

            stream.on('end', () => {
                res.end();
            });

            stream.on('error', (error) => {
                console.error('Stream error:', error);
                res.end();
            });
        } else {
            const response = await apiHandler.generateCompletion(req.body);
            res.json(response);
        }
    } catch (error) {
        console.error('Error:', error);
        res.status(500).json({ error: error.message });
    }
});

app.listen(PORT, () => {
    console.log(`Server running on port ${PORT}`);
});