index.js

const express = require('express');
const cors = require('cors');
const puppeteer = require('puppeteer');
const ffmpeg = require('fluent-ffmpeg');
const ffmpegStatic = require('ffmpeg-static');
const fs = require('fs');
const path = require('path');
const os = require('os');
const { randomUUID } = require('crypto');
const { execSync } = require('child_process');

function hasGpu() {
  try {
    execSync('nvidia-smi', { timeout: 3000, stdio: 'ignore' });
    return true;
  } catch (e) {
    return false;
  }
}

const SERVER_HAS_GPU = hasGpu();

ffmpeg.setFfmpegPath(ffmpegStatic);

const app = express();
app.use(express.json({ limit: '50mb' }));

const requireAuth = (req, res, next) => {
  const apiKey = process.env.STUDIO_API_KEY;
  if (!apiKey) return next();
  const authHeader = req.headers.authorization;
  if (!authHeader || authHeader !== `Bearer ${apiKey}`) {
    console.warn(`[Security] Blocked unauthorized request from IP: ${req.ip}`);
    return res.status(401).json({ error: 'Unauthorized: Invalid or missing API Key' });
  }
  next();
};

const TIME_SHIM = `
  window.__VTIME_INSTALLED__ = true;
  var vtime = 0;
  
  window.__nativeSetTimeout = window.setTimeout;
  window.__nativeClearTimeout = window.clearTimeout;

  const __originalGetContext = HTMLCanvasElement.prototype.getContext;
  HTMLCanvasElement.prototype.getContext = function(type, options) {
    if (type === 'webgl' || type === 'webgl2' || type === 'experimental-webgl') {
      options = options || {};
      options.preserveDrawingBuffer = true;
    }
    if (type === '2d') {
      options = options || {};
      options.willReadFrequently = true;
    }
    return __originalGetContext.call(this, type, options);
  };

  window.__nativeSetInterval = window.setInterval;
  window.__nativeClearInterval = window.clearInterval;
  window.__nativeDate = window.Date;
  window.__nativePerformance = window.performance;

  var _origDate = Date;
  window.Date = function() {
    if (arguments.length === 0) return new _origDate(vtime);
    return new _origDate(...arguments);
  };
  window.Date.now = function() { return vtime; };
  window.Date.parse = _origDate.parse;
  window.Date.UTC = _origDate.UTC;

  if (window.performance) {
    window.performance.now = function() { return vtime; };
  }

  var rafCallbacks = []; var rafId = 0;
  window.__origRaf = window.requestAnimationFrame;
  window.__origCancelRaf = window.cancelAnimationFrame;
  window.requestAnimationFrame = function(cb){ rafId++; rafCallbacks.push({id: rafId, cb: cb}); return rafId; };
  window.cancelAnimationFrame = function(id){ rafCallbacks = rafCallbacks.filter(function(x){ return x.id !== id; }); };

  var timeoutCallbacks = []; var timeoutId = 0;
  window.setTimeout = function(cb, delay){
    timeoutId++;
    timeoutCallbacks.push({id: timeoutId, cb: typeof cb === 'function' ? cb : function(){}, triggerTime: vtime + (delay || 0)});
    return timeoutId;
  };
  window.clearTimeout = function(id){ timeoutCallbacks = timeoutCallbacks.filter(function(x){ return x.id !== id; }); };

  var intervalCallbacks = []; var intervalId = 0;
  window.setInterval = function(cb, delay){
    intervalId++;
    intervalCallbacks.push({id: intervalId, cb: typeof cb === 'function' ? cb : function(){}, interval: delay || 0, nextTime: vtime + (delay || 0)});
    return intervalId;
  };
  window.clearInterval = function(id){ intervalCallbacks = intervalCallbacks.filter(function(x){ return x.id !== id; }); };

  window.__advanceVTime = function(target){
    var STEP = 8;
    while (vtime < target) {
      vtime = Math.min(vtime + STEP, target);
      var pendingTimeouts = timeoutCallbacks;
      timeoutCallbacks = [];
      for (var k = 0; k < pendingTimeouts.length; k++) {
        if (pendingTimeouts[k].triggerTime <= vtime) {
          try { pendingTimeouts[k].cb(); } catch(e){}
        } else {
          timeoutCallbacks.push(pendingTimeouts[k]);
        }
      }
      for (var l = 0; l < intervalCallbacks.length; l++) {
        if (intervalCallbacks[l].nextTime <= vtime) {
          try { intervalCallbacks[l].cb(); } catch(e){}
          intervalCallbacks[l].nextTime = vtime + intervalCallbacks[l].interval;
        }
      }
      var cbs = rafCallbacks; rafCallbacks = [];
      for (var j = 0; j < cbs.length; j++) { try { cbs[j].cb(vtime); } catch(e){} }
      
      if (document.getAnimations) {
        document.getAnimations().forEach(function(a){ 
          try { 
            a.pause(); // CRITICAL: Freeze CSS animations so they don't bleed during page.screenshot() delay
            if (a.__vstartTime === undefined) a.__vstartTime = vtime;
            a.currentTime = vtime - a.__vstartTime; 
          } catch(e){} 
        });
      }
      try {
        if (window.gsap && window.gsap.globalTimeline) {
          window.gsap.globalTimeline.pause();
          window.gsap.globalTimeline.seek(vtime / 1000);
        }
        if (window.__animeInstances) {
          window.__animeInstances.forEach(function(a){ try { a.pause(); a.seek(vtime); } catch(e){} });
        }
      } catch(e){}

    }
  };
`;

function injectShim(html) {
  let processedHtml = html;
  
  const script = `<script>${TIME_SHIM}</script>`;
  if (/<head[^>]*>/i.test(processedHtml)) return processedHtml.replace(/<head[^>]*>/i, (m) => m + script);
  return script + processedHtml;
}

// ----------------------------------------------------------------------------
// ASYNCHRONOUS JOB QUEUE SYSTEM
// ----------------------------------------------------------------------------
const jobs = new Map(); // Store all job states
const queue = [];       // Queue of jobIds waiting to be processed
let isProcessing = false;

async function processQueue() {
  if (isProcessing || queue.length === 0) return;
  isProcessing = true;

  while (queue.length > 0) {
    const jobId = queue.shift();
    const job = jobs.get(jobId);
    if (!job || job.isCancelled) continue;

    // Update positions for everyone left in the queue
    queue.forEach((id, index) => {
      const qJob = jobs.get(id);
      if (qJob) qJob.position = index + 1;
    });

    job.status = 'processing';
    job.position = 0;
    
    try {
      console.log(`[Job ${jobId}] Starting render process...`);
      await executeRenderJob(jobId, job.reqBody);
      job.status = 'done';
    } catch (err) {
      console.error(`[Job ${jobId}] Failed:`, err);
      job.status = 'error';
      job.error = err.message;
      if (job.tempDir) {
        try { fs.rmSync(job.tempDir, { recursive: true, force: true }); } catch (e) {}
      }
    }
  }

  isProcessing = false;
}

async function executeRenderJob(jobId, reqBody) {
  const { html, fps = 60, duration = 5, width = 1920, height = 1080, zoom = 1, panX, panY, exportScale = 1 } = reqBody;
  const tempDir = path.join(os.tmpdir(), `reelify_${jobId}`);
  fs.mkdirSync(tempDir, { recursive: true });

  const totalFrames = Math.round(fps * duration);
  const frameIntervalMs = 1000 / fps;
  const pX = panX ?? (width / 2);
  const pY = panY ?? (height / 2);
  const finalScale = Number((zoom * exportScale).toFixed(3));
  
  const jobState = jobs.get(jobId);
  jobState.tempDir = tempDir;
  
  let totalFramesRendered = 0;

  async function renderFrameChunk(workerId, startFrame, endFrame) {
    let browser;
    try {
      const puppeteerArgs = [
        '--no-sandbox', 
        '--disable-setuid-sandbox', 
        '--disable-dev-shm-usage', 
        '--disable-web-security',
        '--disable-features=IsolateOrigins,site-per-process',
        '--js-flags="--max-old-space-size=4096"'
      ];
      if (!SERVER_HAS_GPU) puppeteerArgs.push('--disable-gpu');

      browser = await puppeteer.launch({
        executablePath: process.env.PUPPETEER_EXECUTABLE_PATH || undefined,
        headless: 'new',
        protocolTimeout: 3600000,
        args: puppeteerArgs
      });
      
      const page = await browser.newPage();
      await page.setViewport({ width, height, deviceScaleFactor: finalScale });
      
      const processedHtml = injectShim(html);
      try {
        await page.setContent(processedHtml, { waitUntil: 'networkidle0', timeout: 30000 });
      } catch (e) {
        console.warn(`[Job ${jobId}] Timeout waiting for networkidle0. Proceeding.`);
      }
      
      const clipWidth = Math.round(width / zoom);
      const clipHeight = Math.round(height / zoom);
      const clipX = Math.round(pX - (clipWidth / 2));
      const clipY = Math.round(pY - (clipHeight / 2));

      // SHARDING FIX: We ALWAYS loop from 0 to endFrame! 
      // This guarantees physics engines evaluate perfectly in order.
      for (let i = 0; i <= endFrame; i++) {
        if (jobState.isCancelled) throw new Error('Job cancelled');
        
        const timeMs = i * frameIntervalMs;
        await page.evaluate((t) => {
          if (typeof window.__renderFrame === 'function') {
            return window.__renderFrame(t);
          } else if (typeof window.__advanceVTime === 'function') {
            window.__advanceVTime(t);
          }
        }, timeMs);

        // If we haven't reached our chunk, fast-forward
        if (i < startFrame) {
          await new Promise(r => setTimeout(r, 5));
          continue;
        }

        await new Promise(r => setTimeout(r, 60));

        const framePath = path.join(tempDir, `frame_${String(i).padStart(5, '0')}.png`);

                await page.evaluate(async (hasGpu) => {
          if (hasGpu) return; 
          const canvases = Array.from(document.querySelectorAll('canvas'));
          await Promise.all(canvases.map((c, index) => {
            return new Promise((resolve) => {
              if (c.width === 0 || c.height === 0) return resolve();
              const dataUrl = c.toDataURL('image/png'); 
              const imgId = 'puppeteer-canvas-composite-' + index;
              let img = document.getElementById(imgId);
              const computedStyle = window.getComputedStyle(c);
              if (!img) {
                img = document.createElement('img');
                img.id = imgId;
                img.style.position = 'fixed';
                img.style.pointerEvents = 'none';
                img.style.zIndex = computedStyle.zIndex || '0';
                c.parentNode.insertBefore(img, c.nextSibling);
              }
              img.onload = resolve;
              img.onerror = resolve;
              img.src = dataUrl;
              const rect = c.getBoundingClientRect();
              img.style.left = rect.left + 'px';
              img.style.top = rect.top + 'px';
              img.style.width = rect.width + 'px';
              img.style.height = rect.height + 'px';
              img.style.transform = computedStyle.transform !== 'none' ? computedStyle.transform : '';
              img.style.opacity = computedStyle.opacity;
              c.style.visibility = 'hidden'; 
            });
          }));
        }, SERVER_HAS_GPU);

        await page.screenshot({ 
          path: framePath,
          type: 'png',
          timeout: 0,
          clip: { x: clipX, y: clipY, width: clipWidth, height: clipHeight }
        });
        
        if (i % 2 === 0) await new Promise(r => setTimeout(r, 100));
        
        totalFramesRendered++;
        jobState.progress = { current: totalFramesRendered, total: totalFrames };
        
        if (totalFramesRendered % 50 === 0) {
          console.log(`[Job ${jobId}] Progress: ${totalFramesRendered}/${totalFrames} frames...`);
        }
      }
      await browser.close();
    } catch (err) {
      if (browser) await browser.close();
      throw err;
    }
  }

  try {
    console.log(`[Job ${jobId}] Extracting frames across 4 multi-core workers...`);
    const NUM_WORKERS = 4;
    const workerPromises = [];
    const framesPerWorker = Math.ceil(totalFrames / NUM_WORKERS);

    for (let w = 0; w < NUM_WORKERS; w++) {
      const startFrame = w * framesPerWorker;
      const endFrame = Math.min((w + 1) * framesPerWorker - 1, totalFrames - 1);
      if (startFrame <= totalFrames - 1) {
        workerPromises.push(renderFrameChunk(w, startFrame, endFrame));
      }
    }

    // Run all 4 browser workers simultaneously!
    await Promise.all(workerPromises);

    const outputPath = path.join(tempDir, 'output.mp4');
    jobState.outputPath = outputPath;

    const isHighRes = (width * exportScale) > 1920;
    const crf = isHighRes ? '23' : '18';

    await new Promise((resolve, reject) => {
      if (jobState.isCancelled) return reject(new Error('Job cancelled'));

      console.log(`[Job ${jobId}] Encoding MP4...`);
      jobState.ffmpegCommand = ffmpeg()
        .input(path.join(tempDir, 'frame_%05d.png'))
        .inputOptions([`-framerate ${fps}`])
        .outputOptions([
          '-c:v libx264',
          '-pix_fmt yuv420p',
          '-preset medium',
          `-crf ${crf}`,
          '-profile:v high',
          '-level 5.2',
          '-vf', 'scale=trunc(iw/2)*2:trunc(ih/2)*2'
        ])
        .output(outputPath)
        .on('end', resolve)
        .on('error', (err) => reject(new Error('FFmpeg encoding failed: ' + err.message)));
        
      jobState.ffmpegCommand.run();
    });

  } catch (err) {
    throw err;
  }
}

// ----------------------------------------------------------------------------
// ENDPOINTS
// ----------------------------------------------------------------------------

app.post('/api/render-async', requireAuth, (req, res) => {
  if (!req.body.html) return res.status(400).json({ error: 'Missing HTML' });
  const jobId = req.body.jobId || randomUUID();
  
  jobs.set(jobId, {
    createdAt: Date.now(),
    id: jobId,
    reqBody: req.body,
    status: 'queued',
    position: queue.length + 1,
    isCancelled: false,
    progress: { current: 0, total: 0 }
  });
  
  queue.push(jobId);
  res.json({ jobId, status: 'queued', position: queue.length });
  
  // Kick off the queue without blocking the response
  processQueue().catch(console.error);
});

app.get('/api/status/:jobId', requireAuth, (req, res) => {
  const job = jobs.get(req.params.jobId);
  if (!job) return res.status(404).json({ error: 'Job not found' });
  
  res.json({
    status: job.status,
    position: job.position,
    progress: job.progress,
    error: job.error
  });
});

app.get('/api/download/:jobId', (req, res) => { // Removed requireAuth for easy browser download
  const job = jobs.get(req.params.jobId);
  if (!job || job.status !== 'done' || !job.outputPath) {
    return res.status(404).json({ error: 'File not ready or not found' });
  }

  console.log(`[Job ${job.id}] Sending it...`);
  res.download(job.outputPath, 'render.mp4', () => {
    // Cleanup after download
    jobs.delete(job.id);
    if (job.tempDir) {
      try { fs.rmSync(job.tempDir, { recursive: true, force: true }); } catch (e) {}
    }
  });
});

app.post('/api/cancel/:jobId', requireAuth, (req, res) => {
  const job = jobs.get(req.params.jobId);
  if (job) {
    job.isCancelled = true;
    if (job.ffmpegCommand) try { job.ffmpegCommand.kill('SIGKILL'); } catch(e){}
    
    const qIndex = queue.indexOf(job.id);
    if (qIndex !== -1) queue.splice(qIndex, 1); // Remove from queue if not started
    
    // Cleanup immediately
    if (job.tempDir) {
      try { fs.rmSync(job.tempDir, { recursive: true, force: true }); } catch (e) {}
    }
    jobs.delete(job.id);
  }
  res.json({ success: true });
});

// Keep original synchronous endpoint for backwards compatibility (local testing)
app.post('/api/render', requireAuth, async (req, res) => {
  // Mock async flow for the sync endpoint
  const jobId = req.body.jobId || randomUUID();
  jobs.set(jobId, { id: jobId, reqBody: req.body, status: 'processing', isCancelled: false });
  try {
    await executeRenderJob(jobId, req.body);
    const job = jobs.get(jobId);
    console.log(`[Job ${job.id}] Sending it...`);
  res.download(job.outputPath, 'render.mp4', () => {
      jobs.delete(jobId);
      if (job.tempDir) try { fs.rmSync(job.tempDir, { recursive: true, force: true }); } catch (e) {}
    });
  } catch (err) {
    res.status(500).json({ error: err.message });
  }
});

const PORT = process.env.PORT || 7860;
const server = app.listen(PORT, () => {
  console.log(`Studio Backend listening on port ${PORT}`);
});
server.setTimeout(0);

// Reaper: Clean up orphaned jobs from memory and disk after 45 minutes
setInterval(() => {
  const now = Date.now();
  for (const [id, job] of jobs.entries()) {
    if (now - job.createdAt > 45 * 60 * 1000) {
      console.log(`[Reaper] Reaping stale job ${id}`);
      if (job.ffmpegCommand) try { job.ffmpegCommand.kill('SIGKILL'); } catch(e){}
      if (job.tempDir) try { fs.rmSync(job.tempDir, { recursive: true, force: true }); } catch(e){}
      jobs.delete(id);
    }
  }
}, 5 * 60 * 1000);