#!/usr/bin/env node

const { spawn } = require('child_process');
const http = require('http');
const net = require('net');
const fs = require('fs');
const path = require('path');

const TTYD_PORT = 4001;
const TTYD_HOST = '0.0.0.0';
const IDLE_TIMEOUT_MS = 5 * 60 * 1000; // 5 minutes of inactivity
const STARTUP_TIMEOUT_MS = 10000; // 10 seconds to start ttyd

let ttydProcess = null;
let lastActivityTime = Date.now();
let clientCount = 0;
let startupPromise = null;

const logFile = '/var/log/ttyd-proxy.log';

function log(message) {
  const timestamp = new Date().toISOString();
  const logLine = `[${timestamp}] ${message}\n`;
  fs.appendFileSync(logFile, logLine);
  console.log(logLine.trim());
}

function startTtyd() {
  if (startupPromise) {
    return startupPromise;
  }

  if (ttydProcess) {
    return Promise.resolve();
  }

  log('Starting ttyd process...');

  startupPromise = new Promise((resolve, reject) => {
    const args = [
      '-W',
      '-p', '4002', // Use port 4002 for the actual ttyd instance
      '-i', '0.0.0.0',
      '/usr/bin/pwsh'
    ];

    // Add password protection if ACCESS_PASSWORD is set
    if (process.env.ACCESS_PASSWORD) {
      args.splice(0, 0, '-c', `user:${process.env.ACCESS_PASSWORD}`);
    }

    ttydProcess = spawn('/usr/local/bin/ttyd', args, {
      env: process.env,
      stdio: ['ignore', 'pipe', 'pipe']
    });

    const startupTimer = setTimeout(() => {
      if (ttydProcess) {
        log('ttyd startup timeout, killing process');
        ttydProcess.kill('SIGKILL');
        ttydProcess = null;
        startupPromise = null;
        reject(new Error('ttyd startup timeout'));
      }
    }, STARTUP_TIMEOUT_MS);

    ttydProcess.stdout.on('data', (data) => {
      log(`ttyd stdout: ${data.slice(0, 200)}`);
    });

    ttydProcess.stderr.on('data', (data) => {
      const stderr = data.toString();
      if (!stderr.includes('client disconnected') && !stderr.includes('new client')) {
        log(`ttyd stderr: ${stderr.slice(0, 200)}`);
      }
    });

    ttydProcess.on('exit', (code, signal) => {
      log(`ttyd exited: code=${code}, signal=${signal}`);
      ttydProcess = null;
      startupPromise = null;
    });

    // Wait for ttyd to start listening on port 4002
    let attempts = 0;
    const checkReady = () => {
      attempts++;
      const client = net.connect(4002, '0.0.0.0', () => {
        clearTimeout(startupTimer);
        client.destroy();
        log('ttyd is ready on port 4002');
        resolve();
      });

      client.on('error', () => {
        client.destroy();
        if (attempts < 50) {
          setTimeout(checkReady, 200);
        } else {
          clearTimeout(startupTimer);
          if (ttydProcess) {
            ttydProcess.kill('SIGKILL');
            ttydProcess = null;
          }
          startupPromise = null;
          reject(new Error('ttyd failed to start'));
        }
      });
    };

    checkReady();
  });

  return startupPromise;
}

function stopTtyd() {
  if (ttydProcess) {
    log('Stopping idle ttyd process...');
    ttydProcess.kill('SIGTERM');
    setTimeout(() => {
      if (ttydProcess) {
        ttydProcess.kill('SIGKILL');
        ttydProcess = null;
      }
    }, 5000);
  }
}

function checkIdleTimeout() {
  const now = Date.now();
  const idleTime = now - lastActivityTime;

  if (clientCount === 0 && idleTime > IDLE_TIMEOUT_MS) {
    log(`ttyd idle for ${idleTime}ms, stopping...`);
    stopTtyd();
  }
}

setInterval(checkIdleTimeout, 60000); // Check every minute

const proxy = http.createServer((clientReq, clientRes) => {
  lastActivityTime = Date.now();

  // Update client count based on request type
  if (clientReq.url === '/' && clientReq.method === 'GET') {
    clientCount++;
  }

  const handleProxy = async () => {
    try {
      await startTtyd();

      const options = {
        hostname: 'localhost',
        port: 4002,
        path: clientReq.url,
        method: clientReq.method,
        headers: {
          ...clientReq.headers,
          host: 'localhost:4002'
        }
      };

      const proxyReq = http.request(options, (proxyRes) => {
        clientRes.writeHead(proxyRes.statusCode, proxyRes.headers);
        proxyRes.pipe(clientRes);

        proxyRes.on('end', () => {
          clientCount = Math.max(0, clientCount - 1);
          lastActivityTime = Date.now();
        });
      });

      proxyReq.on('error', (err) => {
        log(`Proxy request error: ${err.message}`);
        clientCount = Math.max(0, clientCount - 1);
        if (!clientRes.headersSent) {
          clientRes.writeHead(502, { 'Content-Type': 'text/plain' });
          clientRes.end('Bad Gateway: ttyd not available');
        }
      });

      clientReq.pipe(proxyReq);
    } catch (err) {
      log(`Proxy error: ${err.message}`);
      clientCount = Math.max(0, clientCount - 1);
      if (!clientRes.headersSent) {
        clientRes.writeHead(503, { 'Content-Type': 'text/plain' });
        clientRes.end('Service Unavailable: ttyd failed to start');
      }
    }
  };

  handleProxy();
});

proxy.on('upgrade', (clientReq, clientSocket, head) => {
  lastActivityTime = Date.now();
  clientCount++;

  const handleUpgrade = async () => {
    try {
      await startTtyd();

      const proxyReq = http.request({
        hostname: 'localhost',
        port: 4002,
        path: clientReq.url,
        method: clientReq.method,
        headers: {
          ...clientReq.headers,
          host: 'localhost:4002'
        }
      });

      proxyReq.on('upgrade', (proxyRes, proxySocket, proxyHead) => {
        clientSocket.write(
          'HTTP/1.1 101 Switching Protocols\r\n' +
          `Upgrade: ${proxyRes.headers.upgrade}\r\n` +
          `Connection: Upgrade\r\n` +
          Object.entries(proxyRes.headers)
            .filter(([k]) => k !== 'upgrade' && k !== 'connection')
            .map(([k, v]) => `${k}: ${v}`)
            .join('\r\n') +
          '\r\n\r\n'
        );

        proxySocket.pipe(clientSocket).pipe(proxySocket);

        proxySocket.on('close', () => {
          clientCount = Math.max(0, clientCount - 1);
          lastActivityTime = Date.now();
        });

        clientSocket.on('close', () => {
          clientCount = Math.max(0, clientCount - 1);
          lastActivityTime = Date.now();
          proxySocket.end();
        });
      });

      proxyReq.on('error', (err) => {
        log(`WebSocket proxy error: ${err.message}`);
        clientCount = Math.max(0, clientCount - 1);
        clientSocket.end();
      });

      proxyReq.end();
    } catch (err) {
      log(`WebSocket upgrade error: ${err.message}`);
      clientCount = Math.max(0, clientCount - 1);
      clientSocket.end();
    }
  };

  handleUpgrade();
});

proxy.listen(TTYD_PORT, TTYD_HOST, () => {
  log(`ttyd proxy listening on ${TTYD_HOST}:${TTYD_PORT}`);
  log(`ttyd will start on-demand after ${IDLE_TIMEOUT_MS/1000} minutes of inactivity`);
});

process.on('SIGTERM', () => {
  log('Received SIGTERM, shutting down...');
  stopTtyd();
  proxy.close(() => {
    log('Proxy server closed');
    process.exit(0);
  });
});

process.on('SIGINT', () => {
  log('Received SIGINT, shutting down...');
  stopTtyd();
  proxy.close(() => {
    log('Proxy server closed');
    process.exit(0);
  });
});