const http = require('http');
const fs = require('fs');
const WebSocket = require('ws');
require('dotenv').config();


// Configuration
const PORT = 9898;
const API_KEY = 'your_litellm_secret_key'; // consider process.env.API_KEY; // Replace with your actual API key
const LLM_MODEL = 'llama3-70b';
const HTML_FILE = 'index.html';

/**
 * Create and configure HTTP server
 * @returns {http.Server} Configured HTTP server
 */
function createHttpServer() {
  return http.createServer(handleHttpRequest);
}

/**
 * Handle HTTP requests by serving the HTML file
 * @param {http.IncomingMessage} req - HTTP request object
 * @param {http.ServerResponse} res - HTTP response object
 */
function handleHttpRequest(req, res) {
  res.writeHead(200, {'Content-Type': 'text/html'});
  
  fs.readFile(HTML_FILE, (err, data) => {
    if (err) {
      console.error('Error reading HTML file:', err);
      res.statusCode = 500;
      res.end('Error getting the file');
      return;
    }
    res.end(data);
  });
}

/**
 * Set up WebSocket server
 * @param {http.Server} server - HTTP server to attach WebSocket to
 * @returns {WebSocket.Server} Configured WebSocket server
 */
function setupWebSocketServer(server) {
  const wss = new WebSocket.Server({ server });
  
  wss.on('connection', handleWebSocketConnection);
  
  return wss;
}

/**
 * Handle new WebSocket connections
 * @param {WebSocket} ws - WebSocket connection
 */
function handleWebSocketConnection(ws) {
  console.log('Client connected');
  
  ws.on('message', (message) => handleIncomingMessage(ws, message));
  
  ws.on('close', () => {
    console.log('Client disconnected');
  });
}

/**
 * Process incoming WebSocket messages
 * @param {WebSocket} ws - WebSocket connection
 * @param {string} message - Received message
 */
function handleIncomingMessage(ws, message) {
  console.log(`Received message => ${message}`);
  queryLLM(ws, message.toString());
}

/**
 * Query the language model API
 * @param {WebSocket} ws - WebSocket connection to send the response
 * @param {string} message - User message to process
 * @returns {Promise<string|null>} The LLM response or null on error
 */
async function queryLLM(ws, message) {
  try {
    console.log(`\n\n\nWhat is the key: ${API_KEY}`);
    const got = (await import('got')).default;
    const response = await got(process.env.LLM_API_ENDPOINT, {
      method: 'POST',
      headers: {
        'content-type': 'application/json',
        'Authorization': `Bearer ${API_KEY}`,
      },
      json: {
        model: LLM_MODEL,
        messages: [
	   {
             "role": "system",
             "content": "Keep response to 4 lines of text."
          },
          { role: 'user', content: message }
        ],
        max_tokens: 1000
      }
    });

    const data = JSON.parse(response.body);

    if (data?.choices?.[0]?.message?.content) {
      const content = data.choices[0].message.content;
      ws.send(content);
      return content;
    } else {
      console.error('Invalid or empty response from LLM API');
      ws.send('Sorry, I received an invalid response from the language model.');
      return null;
    }
  } catch (error) {
    console.error('Error querying LLM:', error);
    ws.send('Sorry, there was an error processing your request.');
    return null;
  }
}

/**
 * Initialize and start the server
 */
function startServer() {
  const server = createHttpServer();
  setupWebSocketServer(server);
  
  server.listen(PORT, () => {
    console.log(`Server listening on port ${PORT}`);
  });
}

// Start the server when this file is run
startServer();