const express = require('express');
const http = require('http');
const proxy = require('express-http-proxy');
const socketIo = require('socket.io');
const moment = require('moment-timezone');

const app = express();
const server = http.createServer(app);
const io = socketIo(server);

const targetUrl = 'https://api.openai.com';
const openaiKey = process.env.OPENAI_KEY;
const port = 7860;
const baseUrl = getExternalUrl(process.env.SPACE_ID);
const allowedModels = [
    'gpt-3.5-turbo-0125',
    'gpt-3.5-turbo',
    'gpt-3.5-turbo-1106',
    'gpt-3.5-turbo-instruct',
    'gpt-3.5-turbo-16k',
    'gpt-3.5-turbo-0613',
    'gpt-3.5-turbo-16k-0613'
];

let requestDetails = [];
let apiKeyDetails = {}; // Track API key usage and first use timestamp

io.on('connection', (socket) => {
    console.log('A user connected to the websocket for live logs.');
});

function logAndEmit(message, includeTotal = false) {
    if (includeTotal) {
        const totalRequests = getTotalRequests();
        message += ` Total requests so far: ${totalRequests}`;
    }
    console.log(message);
    io.emit('log', message);
}

function getTotalRequests() {
    return Object.values(apiKeyDetails).reduce((acc, details) => acc + details.count, 0);
}


function validateModel(req, res, next) {
    const pathSegments = req.path.split('/');
    // Example path: /v1/engines/gpt-3.5-turbo-0125/completions
    const modelRequested = pathSegments[3]; // Adjust the index as necessary based on your routing

    if (!allowedModels.includes(modelRequested)) {
        console.log(`Forbidden model access attempted: ${modelRequested}`);
        return res.status(403).send(`Forbidden: Model ${modelRequested} is not allowed.`);
    }

    console.log(`Model ${modelRequested} access granted.`);
    next();
}


function authenticateApiKey(req, res, next) {
    const receivedApiKey = req.headers['authorization'];
    if (!receivedApiKey) {
        return res.status(401).send('Unauthorized: No API key provided');
    }

    const validApiKeys = (process.env.SECRET_API_KEYS || '').split(',');
    const apiKeyMatch = validApiKeys.find(key => `Bearer ${key.trim()}` === receivedApiKey);
    
    if (!apiKeyMatch) {
        return res.status(401).send('Unauthorized: API key is invalid');
    }

    const displayKey = apiKeyMatch.trim().substring(0, 7); // Only use the first four characters for display

    if (!apiKeyDetails[apiKeyMatch]) {
        apiKeyDetails[apiKeyMatch] = {
            count: 0,
            firstUse: moment()  // Set the first use to now if it's not already set
        };
    }

    // Check if more than a week has passed since the first use
    if (moment().diff(apiKeyDetails[apiKeyMatch].firstUse, 'weeks') >= 1) {
        return res.status(429).send(`API Key ${displayKey}... is blocked after one week of use.`);
    }

    // Increment the request count
    apiKeyDetails[apiKeyMatch].count++;

    const logMessage = `API Key ${displayKey}... used ${apiKeyDetails[apiKeyMatch].count} times so far.`;
    logAndEmit(logMessage, true);
    
    req.apiKey = apiKeyMatch; // Store the API key for the request
    next();
}

app.use('/api', authenticateApiKey, validateModel, proxy(targetUrl, {
    proxyReqOptDecorator: function(proxyReqOpts, srcReq) {
        proxyReqOpts.headers['Authorization'] = 'Bearer ' + openaiKey;
        return proxyReqOpts;
    },
    proxyReqPathResolver: function(req) {
        const newPath = req.originalUrl.replace('/api', ''); // Adjust the path as necessary
        return newPath;
    }
}));

app.get("/", (req, res) => {
    let requestsInfo = requestDetails.map(detail =>
        `Request ${detail.requestNumber} on ${detail.timestamp} from API Key ${detail.apiKey}... with text "${detail.text}"`).join('<br>');
    res.send(`This is your OpenAI Reverse Proxy URL: ${baseUrl}.<br><br>Requests Overview:<br>${requestsInfo}`);
});

app.get('/logs', (req, res) => {
    res.sendFile(__dirname + '/index.html');
});

function getExternalUrl(spaceId) {
  try {
    const [username, spacename] = spaceId.split("/");
    return `https://${username}-${spacename.replace(/_/g, "-")}.hf.space/api/v1`;
  } catch (e) {
    return "Error generating external URL";
  }
}

server.listen(port, () => {
    const message = `Reverse proxy server and WebSocket running on port ${port}`;
    logAndEmit(message, true); // true to include total requests in the log
});