server/api.ts

/**
 * Optional HTTP API for watermark embedding/detection
 *
 * Serves the web UI and provides API endpoints for server-side processing.
 * Used for HuggingFace Space deployment.
 */

import { createServer } from 'node:http';
import { readFile, stat } from 'node:fs/promises';
import { join, extname } from 'node:path';
import { tmpdir } from 'node:os';
import { randomUUID } from 'node:crypto';
import { writeFile, unlink } from 'node:fs/promises';
import { probeVideo, readYuvFrames, createEncoder } from './ffmpeg-io.js';
import { embedWatermark } from '../core/embedder.js';
import { detectWatermark, detectWatermarkMultiFrame } from '../core/detector.js';
import { getPreset } from '../core/presets.js';
import type { PresetName } from '../core/types.js';

const MIME_TYPES: Record<string, string> = {
  '.html': 'text/html',
  '.js': 'application/javascript',
  '.css': 'text/css',
  '.json': 'application/json',
  '.png': 'image/png',
  '.svg': 'image/svg+xml',
  '.woff2': 'font/woff2',
};

const PORT = parseInt(process.env.PORT || '7860', 10);
const STATIC_DIR = process.env.STATIC_DIR || join(import.meta.dirname || '.', '../dist/web');

async function serveStatic(url: string): Promise<{ data: Buffer; contentType: string } | null> {
  const safePath = url.replace(/\.\./g, '').replace(/\/+/g, '/');
  const filePath = join(STATIC_DIR, safePath === '/' ? 'index.html' : safePath);

  try {
    const s = await stat(filePath);
    if (!s.isFile()) return null;
    const data = await readFile(filePath);
    const ext = extname(filePath);
    return { data, contentType: MIME_TYPES[ext] || 'application/octet-stream' };
  } catch {
    return null;
  }
}

const server = createServer(async (req, res) => {
  const url = new URL(req.url || '/', `http://localhost:${PORT}`);

  // CORS + SharedArrayBuffer headers (required for ffmpeg.wasm)
  res.setHeader('Access-Control-Allow-Origin', '*');
  res.setHeader('Access-Control-Allow-Methods', 'GET, POST, OPTIONS');
  res.setHeader('Access-Control-Allow-Headers', 'Content-Type');
  res.setHeader('Cross-Origin-Opener-Policy', 'same-origin');
  res.setHeader('Cross-Origin-Embedder-Policy', 'require-corp');

  if (req.method === 'OPTIONS') {
    res.writeHead(200);
    res.end();
    return;
  }

  // API: Health check
  if (url.pathname === '/api/health') {
    res.writeHead(200, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify({ status: 'ok' }));
    return;
  }

  // API: Embed
  if (url.pathname === '/api/embed' && req.method === 'POST') {
    try {
      const chunks: Buffer[] = [];
      for await (const chunk of req) chunks.push(chunk as Buffer);
      const body = Buffer.concat(chunks);

      // Parse multipart or raw JSON
      const contentType = req.headers['content-type'] || '';

      if (contentType.includes('application/json')) {
        const { videoBase64, key, preset, payload } = JSON.parse(body.toString());

        const videoBuffer = Buffer.from(videoBase64, 'base64');
        const inputPath = join(tmpdir(), `ltmarx-in-${randomUUID()}.mp4`);
        const outputPath = join(tmpdir(), `ltmarx-out-${randomUUID()}.mp4`);

        await writeFile(inputPath, videoBuffer);

        const config = getPreset((preset || 'moderate') as PresetName);
        const payloadBytes = hexToBytes(payload || 'DEADBEEF');
        const info = await probeVideo(inputPath);

        const encoder = createEncoder(outputPath, info.width, info.height, info.fps);
        const ySize = info.width * info.height;
        const uvSize = (info.width / 2) * (info.height / 2);
        let totalPsnr = 0;
        let frameCount = 0;

        for await (const frame of readYuvFrames(inputPath, info.width, info.height)) {
          const result = embedWatermark(frame.y, info.width, info.height, payloadBytes, key, config);
          totalPsnr += result.psnr;

          const yuvFrame = Buffer.alloc(ySize + 2 * uvSize);
          yuvFrame.set(result.yPlane, 0);
          yuvFrame.set(frame.u, ySize);
          yuvFrame.set(frame.v, ySize + uvSize);
          encoder.stdin.write(yuvFrame);
          frameCount++;
        }

        encoder.stdin.end();
        await new Promise<void>((resolve) => encoder.process.on('close', () => resolve()));

        const outputBuffer = await readFile(outputPath);

        // Cleanup temp files
        await unlink(inputPath).catch(() => {});
        await unlink(outputPath).catch(() => {});

        res.writeHead(200, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify({
          videoBase64: outputBuffer.toString('base64'),
          frames: frameCount,
          avgPsnr: totalPsnr / frameCount,
        }));
      } else {
        res.writeHead(400, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify({ error: 'Expected application/json content type' }));
      }
    } catch (e) {
      res.writeHead(500, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({ error: String(e) }));
    }
    return;
  }

  // API: Detect
  if (url.pathname === '/api/detect' && req.method === 'POST') {
    try {
      const chunks: Buffer[] = [];
      for await (const chunk of req) chunks.push(chunk as Buffer);
      const body = JSON.parse(Buffer.concat(chunks).toString());

      const { videoBase64, key, preset, frames: maxFrames } = body;
      const videoBuffer = Buffer.from(videoBase64, 'base64');
      const inputPath = join(tmpdir(), `ltmarx-det-${randomUUID()}.mp4`);

      await writeFile(inputPath, videoBuffer);

      const config = getPreset((preset || 'moderate') as PresetName);
      const info = await probeVideo(inputPath);
      const framesToRead = Math.min(maxFrames || 10, info.totalFrames);

      const yPlanes: Uint8Array[] = [];
      let count = 0;
      for await (const frame of readYuvFrames(inputPath, info.width, info.height)) {
        yPlanes.push(new Uint8Array(frame.y));
        count++;
        if (count >= framesToRead) break;
      }

      await unlink(inputPath).catch(() => {});

      const result = detectWatermarkMultiFrame(yPlanes, info.width, info.height, key, config);

      res.writeHead(200, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({
        detected: result.detected,
        payload: result.payload ? bytesToHex(result.payload) : null,
        confidence: result.confidence,
        tilesDecoded: result.tilesDecoded,
        tilesTotal: result.tilesTotal,
      }));
    } catch (e) {
      res.writeHead(500, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({ error: String(e) }));
    }
    return;
  }

  // Static file serving
  const staticResult = await serveStatic(url.pathname);
  if (staticResult) {
    res.writeHead(200, { 'Content-Type': staticResult.contentType });
    res.end(staticResult.data);
    return;
  }

  // SPA fallback
  const indexResult = await serveStatic('/');
  if (indexResult) {
    res.writeHead(200, { 'Content-Type': 'text/html' });
    res.end(indexResult.data);
    return;
  }

  res.writeHead(404, { 'Content-Type': 'text/plain' });
  res.end('Not Found');
});

server.listen(PORT, () => {
  console.log(`LTMarx server listening on http://localhost:${PORT}`);
});

function hexToBytes(hex: string): Uint8Array {
  const clean = hex.replace(/^0x/, '');
  const padded = clean.length % 2 ? '0' + clean : clean;
  const bytes = new Uint8Array(padded.length / 2);
  for (let i = 0; i < bytes.length; i++) {
    bytes[i] = parseInt(padded.slice(i * 2, i * 2 + 2), 16);
  }
  return bytes;
}

function bytesToHex(bytes: Uint8Array): string {
  return Array.from(bytes).map((b) => b.toString(16).padStart(2, '0')).join('').toUpperCase();
}