/**
 * Copyright (c) 2023 MERCENARIES.AI PTE. LTD.
 * All rights reserved.
 */
//@ts-check
const VERSION = '0.6.0.hf.013.b';

const express = require('express');
const http = require('http');
const session = require('express-session');
const axios = require('axios');
const { spawn } = require('child_process');
const { createProxyMiddleware } = require('http-proxy-middleware');

const app = express();

const OMNITOOL_INSTALL_SCRIPT = './omnitool_init.sh'; // './omnitool_start.sh';
const CONTAINER_HOST = '127.0.0.1';
const OMNI_URL = 'http://127.0.0.1:1688'; // URL of the OMNITOOL service
const PROXY_PORT_OMNITOOL = 4444;
const CONTAINER_PORT_OMNITOOL = 1688;
const OMNITOOL_HEALTH_PING = '/api/v1/mercenaries/ping';
const FAVICON = 'https://github.com/omnitool-ai/omnitool/raw/main/packages/omni-ui/omni-web/public/favicon.jpg'
const OMNITOOL_SPACE_URL = "omnitool-ai-omnitool-on-hf.hf.space";
const HF_SPACE_URL = "https://huggingface.co/spaces/omnitool-ai/omnitool_on_hf";
const HF_SPACE_DUPLICATE_URL = "https://huggingface.co/spaces/omnitool-ai/omnitool_on_hf?duplicate=true";

const DELAY_OMNITOOL_SET_TO_RUNNING = 2000; // 2 seconds
const CHECK_OMNI_INTERVAL = 60000; // 1 minute

// Global variable
global.OMNITOOL_RUNNING = false;
global.OMNITOOL_READY = false;
global.ALREADY_STARTING = false;
global.LOCAL_URL = "";
global.PING_URL = "";
global.PROTOCOL = "";
global.PROXY_STARTED = false;
global.logs = [];
global.OMNITOOL_PROXY = null;
global.CONNECTED_TO_MASTER = false;



console.log(`************ Omnitool Proxy Server v${VERSION} ************`);

// HTML templates and constants
const LOG_CONTAINER_HTML = `<div id="logContainer"></div>`;
const COMMON_STYLES = `
    <style>
        #logContainer {
            height: 400px;
            overflow-y: scroll;
            background-color: black;
            color: lime;
            font-family: 'Courier New', Courier, monospace;
            padding: 10px;
            white-space: pre-wrap;
            border: 1px solid #ddd;
        }
        .highlight-button {
            animation: pulseAnimation 1s infinite;
            background-color: yellow;
            color: black;
            font-weight: bold;
        }
        .highlight-button-green {
            animation: pulseAnimation 1s infinite;
            background-color: green;
            color: white;
            font-weight: bold;
        }      
        .button-like-link {
            display: inline-block;
            background-color: #f0f0f0; /* Button color */
            color: #000; /* Text color */
            padding: 5px 10px;
            text-align: center;
            text-decoration: none;
            border: 1px solid #000;
            border-radius: 5px;
            font-size: 16px;
            cursor: pointer;
        }
        .button-like-link:hover {
            background-color: #e9e9e9;
        }                     
        @keyframes pulseAnimation {
            0% { transform: scale(1); }
            50% { transform: scale(1.1); }
            100% { transform: scale(1); }
        }
    </style>`;


async function startOmnitoolServer()
{
    if (global.ALREADY_STARTING) return;
    global.ALREADY_STARTING = true;

    console.log('Starting Omnitool Server...');
    return new Promise((resolve, reject) =>
    {
        const omnitoolStartProcess = spawn(OMNITOOL_INSTALL_SCRIPT);
        omnitoolStartProcess.stdout.on('data', (data) =>
        {
            global.logs.push(data.toString());
            console.log(`omnitool stdout: ${data}`);

            if (data.toString().includes(`Server has started and is ready to accept connections`))
            {
                console.log('Omnitool server started successfully');
                setOmnitoolRunning(true);
            }

        });

        omnitoolStartProcess.stderr.on('data', (data) =>
        {
            console.error(`omnitool stderr: ${data}`);
            global.logs.push(data.toString());
        });

        omnitoolStartProcess.on('close', (code) =>
        {
            const message = `Omnitool server process exited with code: ${code}`;
            console.log(message);
            global.logs.push(message);

            global.ALREADY_STARTING = false;
            if (code === 0)
            {
                //@ts-ignore
                resolve();
            } else
            {
                reject(`Omnitool server did not start properly.`);
            }
        });
    });
}

function setOmnitoolRunning(set_running)
{
    if (set_running === true)
    {
        if (global.OMNITOOL_READY === false) 
        {
            global.OMNITOOL_READY = true;
            console.log('Omnitool server is READY! ');
            setTimeout(() =>
            {
                console.log('Omnitool server is RUNNING! ');
                global.OMNITOOL_RUNNING = true;
            }, DELAY_OMNITOOL_SET_TO_RUNNING); // Delay by 2 second
        }

    }
    else
    {
        global.OMNITOOL_READY = false;
        global.OMNITOOL_RUNNING = false;
    }
}

// Function to check the status of the external service
async function checkOmnitoolStatus()
{
    console.log("Checking external service");
    try
    {
        //@ts-ignore
        const response = await axios.get('http://127.0.0.1:1688/api/v1/mercenaries/ping');
        if (response.data && response.data.ping === 'pong' && Object.keys(response.data.payload).length === 0)
        {
            setOmnitoolRunning(true);
        }
        else
        {
            setOmnitoolRunning(false);
        }
    } catch (error)
    {
        console.error('Cannot access OMNITOOL. It is probably not running...', error.message);
        setOmnitoolRunning(false);
    }
}

async function handleGetStartOmnitoolServer(req, res)
{
    console.log(`Omnitool Server:ALREADY_STARTING =  ${global.ALREADY_STARTING}`);
    if (global.ALREADY_STARTING) 
    {
        res.writeHead(200, { 'Content-Type': 'text/html' });
        res.end("Omnitool server already starting");
        return;
    }
    try
    {
        await startOmnitoolServer();
        res.writeHead(200, { 'Content-Type': 'text/html' });
        res.end("Omnitool server started successfully");
    }
    catch (error)
    {
        console.error(error);
        global.ALREADY_STARTING = false;
        res.writeHead(500, { 'Content-Type': 'text/html' });
        res.end(`Error starting Omnitool server: ${error}`);
    }
}

async function handleGetOmnitoolLogs(req, res)
{
    res.writeHead(200, { 'Content-Type': 'application/json' });
    const reply = { logs: global.logs, ready: global.OMNITOOL_READY };
    res.end(JSON.stringify(reply));
    return;
}

async function handleGetBurstIframe(req, res)
{
    res.writeHead(200, { 'Content-Type': 'application/json' });
    req.session.isVisited = true;

    const reply = { burst_iframe: true };
    res.end(JSON.stringify(reply));
    return;
}

async function proxyRequest(req, res)
{

    console.log('Proxying request');
    if (global.PROXY_STARTED) return;
    global.PROXY_STARTED = true;

    // Proxy logic...
    const options = { hostname: CONTAINER_HOST, port: CONTAINER_PORT_OMNITOOL, path: req.url, method: req.method, headers: req.headers, };
    const proxy = http.request(options, (proxyRes) =>
    {
        res.writeHead(proxyRes.statusCode, proxyRes.headers);
        proxyRes.pipe(res, { end: true });
    });
    req.pipe(proxy, { end: true });
    return;
}

function getButtonsString(buttonsHTML)
{
    // Message to display based on CONNECTED_TO_MASTER
    let connectionStatusMessage;
    if (global.CONNECTED_TO_MASTER)
    {
        connectionStatusMessage = '<div>You are now connected DIRECTLY to the REFERENCE Omnitool Space on Huggingface.</div><div style="color: red;">It is HIGHLY recommended that you duplicate this space and make it private to secure your keys, recipes and outputs.</div><div>When DUPLICATING the Space, consider choosing the 20 Gig option to persist keys, recipes and outputs between server restarts. Those will still be lost, however, if the Space is rebuilt.</div>';
    } else
    {
        connectionStatusMessage = `<div style="color: green;">You are now connected to the space ${global.LOCAL_URL}.</div><div style="color: red;">If this is NOT yours, it is HIGHLY recommended that you duplicate this space and make it private to secure your keys, recipes and outputs.</div><div>When DUPLICATING the Space, consider choosing the 20 Gig option to persist keys, recipes and outputs between server restarts. Those will still be lost, however, if the Space is rebuilt.</div>`;
    }

    return `
        <html>
            <head>
                <title>Omnitool</title>
                ${COMMON_STYLES}
            </head>
            <body>
                <h1>Omnitool.ai on Hugging Face</h1>
                <p>Version: ${VERSION}</p>
                <div>Welcome to Omnitool.ai running on Hugging Face!</div>
                <div>This is a fully functional version, but all recipes and outputs are lost when the server is restarted!</div>
                <div>For more details, including how to install omnitool on your own computer, visit us at <a href="https://omnitool.ai" target="_blank" class="button-like-link" id="duplicateRepoButton">omnitool.ai</a></div>
                <p>-----------------------------------</p>
                <p>This Space URL: ${global.LOCAL_URL}</p>
                <p>Using REFERENCE Space: ${global.CONNECTED_TO_MASTER}</p>
                <p>Omnitool SERVER Launched: ${global.ALREADY_STARTING}</p>
                <p>Omnitool SERVER Ready: ${global.OMNITOOL_READY}</p>
                <p>-----------------------------------</p>
                <p></p>
                ${connectionStatusMessage}
                <p></p>
                ${buttonsHTML}
                ${LOG_CONTAINER_HTML}
                
                <script>
                    // Script to check Hugging Face login status and control 'Clone Repo' button
                
                    function startServer() {
                        document.getElementById('startServerButton').classList.remove
                        ('highlight-button');
                        document.getElementById('startServerButton').disabled = true;


                        fetch('/start-omnitool-server')
                            .then(response => response.json())
                            .then(data => startLogPolling());
  
                        
                    }

                    function exitIframe() {
                        fetch('/burst-iframe')
                            .then(response => response.json())
                            .then(data => window.open('/', '_blank'));
                    }

                    function startLogPolling() {

                        const interval = setInterval(() => {
                            fetch('/omnitool-logs')
                                .then(response => response.json())
                                .then(data => {
                                    const logContainer = document.getElementById('logContainer');
                                    logContainer.innerText = data.logs.join("");
                                    if (data.ready) {
                                        clearInterval(interval);
                                        document.getElementById('exitIframeButton').disabled = false;
                                        document.getElementById('exitIframeButton').classList.add('highlight-button');


                                    }
                                    else {
                                        scrollToBottom(logContainer);
                                    }
                                });
                        }, 500);
                    }

                    function scrollToBottom(element) {
                        element.scrollTop = element.scrollHeight;
                    }
                    
                    startLogPolling();
                </script>
            </body>
        </html>`;
}


async function handleGetRoot(req, res)
{
    setGlobals(req);
    
    console.log(`req.session.isVisited = ${req.session.isVisited}`);
    console.log(`global.OMNITOOL_RUNNING = ${global.OMNITOOL_RUNNING}`);

    if (!req.session.isVisited || !global.OMNITOOL_RUNNING)
    {
        await checkOmnitoolStatus();
    }

    if (!global.OMNITOOL_RUNNING)
    {
        console.log('Omnitool server is not running');
        let startButtonClass = '';
        if (!global.ALREADY_STARTING) startButtonClass = 'highlight-button';
        const gotoButtonClass = '';

        let buttonsHTML = `
        <button id="startServerButton" class="${startButtonClass}" onclick="startServer()">CREATE Omnitool Server</button>
        <button id="exitIframeButton" class="${gotoButtonClass}" onclick="exitIframe()" disabled>LAUNCH Omnitool</button>
        <a href="${HF_SPACE_DUPLICATE_URL}" target="_blank" class="button-like-link" id="duplicateRepoButton">DUPLICATE Space</a>`;

        const html = getButtonsString(buttonsHTML);
        res.writeHead(200, { 'Content-Type': 'text/html' });
        res.end(html);

        return;
    }

    if (!req.session.isVisited)
    {
        console.log('First time visitor');
        const gotoButtonClass = 'highlight-button-green';

        let buttonsHTML = `
        <button id="exitIframeButton" class="${gotoButtonClass}" onclick="exitIframe()">GOTO OMNITOOL</button>
        <a href="${HF_SPACE_DUPLICATE_URL}" target="_blank" class="button-like-link" id="duplicateRepoButton">DUPLICATE SPACE</a>`;
        const html = getButtonsString(buttonsHTML);
        res.writeHead(200, { 'Content-Type': 'text/html' });
        res.end(html);

        return;
    }

    proxyRequest(req, res);
}

function setGlobals(req)
{


    console.log(`Setting global using req.headers['host'] = ${req.headers['host']}`);
    global.LOCAL_URL = req.headers['host'];

    // !!!!!!!!!
    // TEST ONLY: global.LOCAL_URL = 'manu-sapiens-omnitool-test-3.hf.space/';
    // !!!!!!!!!

    if (req.protocol === "https") global.PROTOCOL = 'https'; else global.PROTOCOL = 'http';
    console.log(`global.LOCAL_URL = ${global.LOCAL_URL}\nglobal.PROTOCOL = ${global.PROTOCOL}`);

    if (!global.LOCAL_URL) throw new Error('No host header found in request' + JSON.stringify(req.headers, null, 2));

    const hostname = global.LOCAL_URL.split(':')[0];
    const newUrl = `${global.PROTOCOL}://${hostname}`;//:${PROXY_PORT_OMNITOOL}`;
    global.PING_URL = `${newUrl}${OMNITOOL_HEALTH_PING}`;
    console.log(`global.PING_URL = ${global.PING_URL}`);

    // New logic to set CONNECTED_TO_MASTER
    global.CONNECTED_TO_MASTER = false;
    if (global.LOCAL_URL === OMNITOOL_SPACE_URL) global.CONNECTED_TO_MASTER = true;
    console.log(`global.CONNECTED_TO_MASTER = ${global.CONNECTED_TO_MASTER}`);

}

async function startServer()
{
    // Start server
    http.createServer(app).listen(PROXY_PORT_OMNITOOL, () =>
    {
        console.log(`Server running on port ${PROXY_PORT_OMNITOOL}`);
    });

}

function omnitoolProxyMiddleware(req, res, next)
{
    if (global.OMNITOOL_RUNNING && req.session.isVisited)
    {
        // Proxy all requests to OMNITOOL when conditions are met
        return createProxyMiddleware({
            target: OMNI_URL,
            changeOrigin: true,
            ws: true // if you need WebSocket support
        })(req, res, next);
    } else
    {
        // Continue with normal processing for other cases
        next();
    }
}


async function main(app)
{
    await checkOmnitoolStatus();

    // Configure session middleware
    app.use(session({
        secret: 'your-secret-key',
        resave: false,
        saveUninitialized: true
    }));
    try
    {
        await startServer();
    }
    catch (error)
    {
        console.error(`There was an error starting the server: ${error}`);
        return;
    }


    app.use(omnitoolProxyMiddleware);

    app.get('/favicon.ico', (req, res) => res.status(204));
    app.get('/start-omnitool-server', (req, res) => handleGetStartOmnitoolServer(req, res));
    app.get('/omnitool-logs', (req, res) => handleGetOmnitoolLogs(req, res));
    app.get('/burst-iframe', (req, res) => handleGetBurstIframe(req, res));
    app.get('/', (req, res) => handleGetRoot(req, res));

    setInterval(async () =>
    {
        await checkOmnitoolStatus();
    }, CHECK_OMNI_INTERVAL);

}

// Define the proxy middleware outside the function
global.OMNITOOL_PROXY = createProxyMiddleware({
    target: OMNI_URL,
    changeOrigin: true,
    ws: true // if you need WebSocket support
});

function omnitoolProxyMiddleware(req, res, next)
{
    if (global.OMNITOOL_RUNNING && req.session.isVisited)
    {
        // Use the predefined proxy middleware
        return global.OMNITOOL_PROXY(req, res, next);
    } else
    {
        // Continue with normal processing for other cases
        next();
    }
}
main(app);