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);

        try {
            if (payload.stream) {
                return this._handleStreamingResponse(payload);
            } else {
                return this._handleRegularResponse(payload);
            }
        } catch (error) {
            throw new Error(`API request failed: ${error.message}`);
        }
    }

    async _handleRegularResponse(payload) {
        const response = await axios.post(this.API_URL, payload, {
            headers: this.headers
        });
        return response.data;
    }

    async* _handleStreamingResponse(payload) {
        const response = await axios.post(this.API_URL, payload, {
            headers: this.headers,
            responseType: 'stream'
        });

        for await (const chunk of response.data) {
            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]') continue;
                        const deltaContent = content?.choices?.[0]?.delta?.content;
                        if (deltaContent) yield deltaContent;
                    } catch (error) {
                        continue;
                    }
                }
            }
        }
    }
}

module.exports = { DeepInfraHandler };