TransferTribe/pages/api/chunk-upload.ts

// pages/api/chunk-upload.ts
// Chunked upload endpoint for large files with encryption

import { NextApiRequest, NextApiResponse } from 'next';
import { IncomingForm } from 'formidable';
import fs from 'fs';
import { promises as fsPromises } from 'fs';
import path from 'path';
import { v4 as uuidv4 } from 'uuid';
import { UPLOAD_DIR } from '@/lib/config';
import { prisma } from '@/lib/prisma';
import bcrypt from 'bcrypt';
import { PLAN_CONFIG } from '@/lib/subscription';
import { sendMailjetEmail } from '@/lib/mailjet';
import { getServerSession } from 'next-auth/next';
import { authOptions } from '@/lib/auth';

export const config = {
  api: {
    bodyParser: false,
  },
};

interface ChunkMetadata {
  uploadId: string;
  chunkIndex: number;
  totalChunks: number;
  chunkSize: number;
}

interface UploadSession {
  uploadId: string;
  totalSize: number;
  totalChunks: number;
  receivedChunks: Set<number>;
  chunkSize: number;
  sender: string;
  recipient: string;
  password?: string;
  encryptionIv: string;
  originalFilename: string;
  message?: string;
  filenames?: string;
  createdAt: Date;
}

// Store active upload sessions in memory (in production, use Redis)
const uploadSessions = new Map<string, UploadSession>();

// Cleanup abandoned uploads every 5 minutes (default 24h timeout)
const SESSION_TIMEOUT = 24 * 60 * 60 * 1000; // 24 hours
setInterval(() => {
  const now = Date.now();
  for (const [uploadId, session] of uploadSessions.entries()) {
    if (now - session.createdAt.getTime() > SESSION_TIMEOUT) {
      uploadSessions.delete(uploadId);
      // Clean up temp files
      const chunkDir = path.join(UPLOAD_DIR, '.tmp', uploadId);
      fsPromises.rm(chunkDir, { recursive: true }).catch(err => 
        console.warn(`Failed to clean abandoned upload ${uploadId}:`, err)
      );
    }
  }
}, 5 * 60 * 1000); // Run every 5 minutes

export default async function handler(req: NextApiRequest, res: NextApiResponse) {
  const session = await getServerSession(req, res, authOptions);
  if (!session?.user?.email) {
    return res.status(401).json({ success: false, message: 'Unauthorized' });
  }
  const sessionEmail = session.user.email;

  if (req.method === 'POST') {
    return handleChunkUpload(req, res, sessionEmail);
  } else if (req.method === 'GET') {
    return handleStatusCheck(req, res, sessionEmail);
  } else if (req.method === 'PUT') {
    return handleChunkComplete(req, res, sessionEmail);
  }

  return res.status(405).json({ success: false, message: 'Method not allowed' });
}

async function handleChunkUpload(
  req: NextApiRequest,
  res: NextApiResponse,
  sessionEmail: string
) {
  const form = new IncomingForm({
    uploadDir: path.join(UPLOAD_DIR, '.tmp'),
    keepExtensions: true,
  });

  // Ensure temp directory exists
  try {
    await fsPromises.mkdir(path.join(UPLOAD_DIR, '.tmp'), { recursive: true });
  } catch (err) {
    console.error('Failed to create temp directory:', err);
  }

  return new Promise((resolve) => {
    form.parse(req, async (err: Error | null, fields: any, files: any) => {
      if (err) {
        return resolve(res.status(400).json({ success: false, message: 'Parse error: ' + err.message }));
      }

      try {
        const chunkFile = Array.isArray(files.chunk) ? files.chunk[0] : files.chunk;
        const uploadId = Array.isArray(fields.uploadId) ? fields.uploadId[0] : fields.uploadId;
        const chunkIndex = parseInt(Array.isArray(fields.chunkIndex) ? fields.chunkIndex[0] : fields.chunkIndex);
        const totalChunks = parseInt(Array.isArray(fields.totalChunks) ? fields.totalChunks[0] : fields.totalChunks);
        const chunkSize = parseInt(Array.isArray(fields.chunkSize) ? fields.chunkSize[0] : fields.chunkSize);

        if (!chunkFile || !uploadId || isNaN(chunkIndex) || isNaN(totalChunks)) {
          return resolve(res.status(400).json({ success: false, message: 'Missing chunk metadata' }));
        }

        // Get or create upload session
        let session = uploadSessions.get(uploadId);
        if (session && session.sender !== sessionEmail) {
          return resolve(res.status(403).json({ success: false, message: 'Forbidden' }));
        }
        if (!session) {
          const sender = sessionEmail;
          const recipient = Array.isArray(fields.email) ? fields.email[0] : fields.email;
          const password = Array.isArray(fields.password) ? fields.password[0] : fields.password;
          const encryptionIv = Array.isArray(fields.encryptionIv) ? fields.encryptionIv[0] : fields.encryptionIv;
          const originalFilename = Array.isArray(fields.originalFilename) ? fields.originalFilename[0] : fields.originalFilename;
          const message = Array.isArray(fields.message) ? fields.message[0] : fields.message;
          const filenames = Array.isArray(fields.filenames) ? fields.filenames[0] : fields.filenames;

          if (!sender || !recipient || !encryptionIv) {
            return resolve(res.status(400).json({ success: false, message: 'Missing required fields' }));
          }

          // Verify user exists
          const user = await prisma.user.findUnique({ where: { email: sender } });
          if (!user) {
            return resolve(res.status(404).json({ success: false, message: 'User not found' }));
          }

          session = {
            uploadId,
            totalSize: chunkSize * totalChunks,
            totalChunks,
            receivedChunks: new Set(),
            chunkSize,
            sender,
            recipient,
            password,
            encryptionIv,
            originalFilename,
            message,
            filenames,
            createdAt: new Date(),
          };

          // Check plan limits before accepting upload
          const plan = (user.plan || 'free') as 'free' | 'rookie' | 'pro';
          const limits = PLAN_CONFIG[plan];

          if (limits.maxFileSize !== Infinity && session.totalSize > limits.maxFileSize) {
            return resolve(res.status(413).json({
              success: false,
              message: `File exceeds max size for your plan (${plan})`
            }));
          }

          uploadSessions.set(uploadId, session);
        }

        // Save chunk to temp location
        const chunkDir = path.join(UPLOAD_DIR, '.tmp', uploadId);
        await fsPromises.mkdir(chunkDir, { recursive: true });
        const chunkPath = path.join(chunkDir, `chunk-${chunkIndex}`);

        await new Promise<void>((resolve, reject) => {
          const readStream = fs.createReadStream(chunkFile.filepath);
          const writeStream = fs.createWriteStream(chunkPath);

          readStream.on('error', reject);
          writeStream.on('error', reject);
          writeStream.on('finish', resolve);

          readStream.pipe(writeStream);
        });

        // Clean up formidable temp file
        try {
          await fsPromises.unlink(chunkFile.filepath);
        } catch (err) {
          console.warn('Failed to clean up temp file:', err);
        }

        session.receivedChunks.add(chunkIndex);

        return resolve(res.status(200).json({
          success: true,
          uploadId,
          chunkIndex,
          receivedChunks: Array.from(session.receivedChunks),
        }));
      } catch (err: any) {
        console.error('Chunk upload error:', err);
        return resolve(res.status(500).json({ success: false, message: 'Upload error: ' + err.message }));
      }
    });
  });
}

async function handleStatusCheck(
  req: NextApiRequest,
  res: NextApiResponse,
  sessionEmail: string
) {
  const { uploadId } = req.query;

  if (!uploadId || typeof uploadId !== 'string') {
    return res.status(400).json({ success: false, message: 'Missing uploadId' });
  }

  const session = uploadSessions.get(uploadId);
  if (!session || session.sender !== sessionEmail) {
    return res.status(404).json({ success: false, message: 'Upload session not found' });
  }

  return res.status(200).json({
    success: true,
    uploadId,
    totalChunks: session.totalChunks,
    receivedChunks: Array.from(session.receivedChunks),
    isComplete: session.receivedChunks.size === session.totalChunks,
  });
}

async function handleChunkComplete(
  req: NextApiRequest,
  res: NextApiResponse,
  sessionEmail: string
) {
  const { uploadId } = req.query;

  if (!uploadId || typeof uploadId !== 'string') {
    return res.status(400).json({ success: false, message: 'Missing uploadId' });
  }

  const session = uploadSessions.get(uploadId);
  if (!session || session.sender !== sessionEmail) {
    return res.status(404).json({ success: false, message: 'Upload session not found' });
  }

  // Check if all chunks received
  if (session.receivedChunks.size !== session.totalChunks) {
    return res.status(400).json({
      success: false,
      message: `Missing chunks. Received ${session.receivedChunks.size}/${session.totalChunks}`,
    });
  }

  try {
    // Reassemble chunks into final file
    const chunkDir = path.join(UPLOAD_DIR, '.tmp', uploadId);
    const finalFilename = uuidv4() + '.enc';
    const finalPath = path.join(UPLOAD_DIR, finalFilename);

    const writeStream = fs.createWriteStream(finalPath);

    // Honour backpressure so a multi-GB reassembly flushes to disk instead of
    // queueing in the stream's internal buffer.
    const write = (buf: Buffer): Promise<void> =>
      new Promise((resolve, reject) => {
        if (writeStream.write(buf)) return resolve();
        writeStream.once('drain', resolve);
        writeStream.once('error', reject);
      });

    for (let i = 0; i < session.totalChunks; i++) {
      const chunkPath = path.join(chunkDir, `chunk-${i}`);
      // Retry reading with exponential backoff to handle file locks
      let chunkData: Buffer | null = null;
      for (let attempt = 0; attempt < 3; attempt++) {
        try {
          chunkData = await fsPromises.readFile(chunkPath);
          break;
        } catch (err: any) {
          if (attempt < 2 && err.code === 'EACCES') {
            await new Promise(resolve => setTimeout(resolve, Math.pow(2, attempt) * 100));
          } else {
            throw err;
          }
        }
      }
      // A silently skipped chunk would corrupt the file, so fail loudly.
      if (!chunkData) {
        throw new Error(`Chunk ${i} missing during reassembly`);
      }
      await write(chunkData);
    }

    await new Promise<void>((resolve, reject) => {
      writeStream.on('finish', resolve);
      writeStream.on('error', reject);
      writeStream.end();
    });

    // Get final file size
    const stats = await fsPromises.stat(finalPath);

    // Verify user and check monthly limits
    const user = await prisma.user.findUnique({ where: { email: session.sender } });
    if (!user) {
      await fsPromises.unlink(finalPath);
      return res.status(404).json({ success: false, message: 'User not found' });
    }

    const plan = (user.plan || 'free') as 'free' | 'rookie' | 'pro';
    const limits = PLAN_CONFIG[plan];

    const startOfMonth = new Date();
    startOfMonth.setDate(1);
    startOfMonth.setHours(0, 0, 0, 0);

    const transfersThisMonth = await prisma.transfer.findMany({
      where: {
        senderEmail: session.sender,
        createdAt: { gte: startOfMonth },
      },
    });

    const totalSizeSentThisMonth = transfersThisMonth.reduce((acc, t) => acc + BigInt(t.totalSize), BigInt(0));
    if (
      (limits.maxTransfersPerMonth !== Infinity && transfersThisMonth.length >= limits.maxTransfersPerMonth) ||
      (limits.maxTransferSizePerMonth !== Infinity && totalSizeSentThisMonth + BigInt(stats.size) > BigInt(limits.maxTransferSizePerMonth))
    ) {
      await fsPromises.unlink(finalPath);
      return res.status(429).json({ success: false, message: 'Plan limits exceeded' });
    }

    // Hash password
    const hash = session.password ? await bcrypt.hash(session.password, 10) : null;

    // Parse filenames
    let totalFiles = 1;
    if (session.filenames) {
      try {
        const filesArray = JSON.parse(session.filenames);
        if (Array.isArray(filesArray)) {
          totalFiles = filesArray.length;
        }
      } catch {}
    }

    const expiresAt = new Date(Date.now() + limits.maxExpiryMs);

    // Create transfer record
    const transfer = await prisma.transfer.create({
      data: {
        senderEmail: session.sender,
        recipientEmail: session.recipient,
        passwordHash: hash,
        downloadUrl: uuidv4(),
        expiresAt,
        filenames: session.filenames || null,
        message: session.message || null,
        totalFiles,
        totalSize: stats.size,
        userId: user.id,
      },
    });

    // Create file record separately
    await prisma.file.create({
      data: {
        name: session.originalFilename || 'file',
        path: finalPath,
        size: stats.size,
        transferId: transfer.id,
      },
    });

    // Update transfer with encryption IV in filenames field (temporary workaround)
    // TODO: Fix Prisma type generation for encryptionIv field
    if (session.encryptionIv) {
      // Store IV in a separate metadata approach or update via raw query
      await (prisma as any).$executeRaw`
        UPDATE "File" SET "encryptionIv" = ${session.encryptionIv}
        WHERE "transferId" = ${transfer.id}
      `;
    }

    // Send email
    const link = `${process.env.NEXT_PUBLIC_APP_URL}/download/${transfer.downloadUrl}`;

    await sendMailjetEmail({
      to: session.recipient,
      subject: 'You\'ve received an encrypted file',
      text: `
${session.sender} sent you an encrypted file via TransferTribe.

Link: ${link}

${session.message ? `Message:\n${session.message}\n\n` : ''}Note: You'll need the password they shared with you to decrypt it.
`.trim(),
    });

    // Clean up session and temp files
    uploadSessions.delete(uploadId);
    try {
      await fsPromises.rm(chunkDir, { recursive: true });
    } catch (err) {
      console.warn('Failed to clean up temp directory:', err);
    }

    return res.status(200).json({
      success: true,
      message: 'Upload complete',
      downloadUrl: transfer.downloadUrl,
    });
  } catch (err: any) {
    console.error('Chunk completion error:', err);
    return res.status(500).json({ success: false, message: 'Completion error: ' + err.message });
  }
}