import fs from 'fs';
import dotenv from 'dotenv';
import { logger } from './logger';

dotenv.config();

const LMSTUDIO_BASE_URL = process.env.LMSTUDIO_BASE_URL;
const LMSTUDIO_API_KEY = process.env.LMSTUDIO_API_KEY;
const LMSTUDIO_MODEL = process.env.LMSTUDIO_MODEL;

async function callLmstudio(prompt: string): Promise<any> {
    if (!LMSTUDIO_BASE_URL) {
        throw new Error('LMSTUDIO_BASE_URL is not defined in the .env file');
    }

    for (let i = 0; i < 10; i++) {
        let llmResponse = "";

        try {
            const response = await fetch(`${LMSTUDIO_BASE_URL}/chat/completions`, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                    'Authorization': `Bearer ${LMSTUDIO_API_KEY}`,
                },
                body: JSON.stringify({
                    model: LMSTUDIO_MODEL,
                    messages: [
                        {
                            role: 'user',
                            content: prompt,
                        },
                    ],
                    temperature: 0.7,
                }),
            });

            const data = await response.json();
            if (data.choices && data.choices.length > 0) {
                const content = data.choices[0].message.content;
                llmResponse = content;
                const jsonMatch = content.match(/\{[\s\S]*\}/);
                if (jsonMatch) {
                    return JSON.parse(jsonMatch[0]);
                } else {
                    const arrayMatch = content.match(/\[[\s\S]*\]/);
                    if (arrayMatch) {
                        return JSON.parse(arrayMatch[0]);
                    }
                }
                // If no JSON/array found, return the raw content
                return content;
            } else {
                logger.error('Unexpected API response:', data);
            }
        } catch (error) {
            logger.error(`Attempt ${i + 1} failed:`, error);
            logger.debug(`LLM response: ${llmResponse}`)
        }
    }

    throw new Error('Failed to get response from LLM after 10 attempts');
}

async function callLMStudioAPIWithFile(imagePath: string, prompt: string): Promise<any> {
    if (!LMSTUDIO_BASE_URL) {
        throw new Error('LMSTUDIO_BASE_URL is not defined in the .env file');
    }

    const imageBuffer = fs.readFileSync(imagePath);
    const base64Image = imageBuffer.toString('base64');

    for (let i = 0; i < 10; i++) {
        let llmResponse = "";

        try {
            const response = await fetch(`${LMSTUDIO_BASE_URL}/chat/completions`, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                    'Authorization': `Bearer ${LMSTUDIO_API_KEY}`,
                },
                body: JSON.stringify({
                    model: LMSTUDIO_MODEL,
                    messages: [
                        {
                            role: 'user',
                            content: [
                                { type: 'image_url', image_url: { url: `data:image/jpeg;base64,${base64Image}` } },
                                { type: 'text', text: prompt },
                            ],
                        },
                    ],
                    temperature: 0.7,
                }),
            });

            const data = await response.json();
            if (data.choices && data.choices.length > 0) {
                const content = data.choices[0].message.content;
                llmResponse = content;
                const jsonMatch = content.match(/\{[\s\S]*\}/);
                if (jsonMatch) {
                    return JSON.parse(jsonMatch[0]);
                } else {
                    const arrayMatch = content.match(/\[[\s\S]*\]/);
                    if (arrayMatch) {
                        return JSON.parse(arrayMatch[0]);
                    }
                }
                return content;
            } else {
                logger.error('Unexpected API response:', data);
            }
        } catch (error) {
            logger.error(`Attempt ${i + 1} failed:`, error);
            logger.debug(`LLM response: ${llmResponse}`)
        }
    }

    throw new Error('Failed to describe image after 10 attempts');
}

export { callLmstudio, callLMStudioAPIWithFile };
