Code.js

// --- CONFIGURATION ---
const props = PropertiesService.getScriptProperties();
const CONFIG = {
  MATTER_ID: props.getProperty("MATTER_ID"),
  TARGET_USER: props.getProperty("TARGET_USER"),
  
  // Folder ID for the extracted MBOX file
  MBOX_FOLDER_ID: props.getProperty("MBOX_FOLDER_ID")
};

/**
 * Main entry point: Initiates the exports and sets up the first trigger.
 */
function initiateExport() {
  const TOKEN = ScriptApp.getOAuthToken();
  const jobs = [
    { type: "28day", status: "PENDING", exportId: null },
    { type: "365day", status: "PENDING", exportId: null }
  ];

  try {
    const exportUrl = `https://vault.googleapis.com/v1/matters/${CONFIG.MATTER_ID}/exports`;
    const now = new Date();
    const twentyEightDaysAgo = new Date();
    twentyEightDaysAgo.setDate(now.getDate() - 28);
    const threeSixtyFiveDaysAgo = new Date();
    threeSixtyFiveDaysAgo.setDate(now.getDate() - 365);

    // 1. Initiate Exports
    jobs.forEach(job => {
      Logger.log(`Initiating ${job.type} export for User: ${CONFIG.TARGET_USER}...`);
      
      const payload = {
        name: `Gmail Export (${job.type}) - ` + now.toISOString(),
        query: {
          corpus: "MAIL",
          dataScope: "ALL_DATA",
          searchMethod: "ACCOUNT",
          accountInfo: { emails: [CONFIG.TARGET_USER] }
        },
        exportOptions: { mailOptions: { exportFormat: "MBOX" } }
      };

      if (job.type === "28day") {
        payload.query.startTime = twentyEightDaysAgo.toISOString();
        payload.query.endTime = now.toISOString();
      } else if (job.type === "365day") {
        payload.query.startTime = threeSixtyFiveDaysAgo.toISOString();
        payload.query.endTime = now.toISOString();
      }

      const createResponse = UrlFetchApp.fetch(exportUrl, {
        method: "post",
        contentType: "application/json",
        headers: { Authorization: `Bearer ${TOKEN}` },
        payload: JSON.stringify(payload),
        muteHttpExceptions: true
      });

      if (createResponse.getResponseCode() !== 200) {
        throw new Error(`Failed to create ${job.type} export. Code: ${createResponse.getResponseCode()} | Response: ${createResponse.getContentText()}`);
      }

      const exportData = JSON.parse(createResponse.getContentText());
      job.exportId = exportData.id;
      job.status = "IN_PROGRESS";
      Logger.log(`${job.type} export initiated. Export ID: ${job.exportId}`);
    });

    // 2. Save state to PropertiesService
    const props = PropertiesService.getScriptProperties();
    props.setProperty("EXPORT_JOBS", JSON.stringify(jobs));
    props.setProperty("POLL_COUNT", "0");

    // 3. Setup first trigger (check in 2 minutes)
    setupTrigger();
    Logger.log("Initial exports triggered and status polling scheduled.");

  } catch (e) {
    Logger.log("CRITICAL ERROR in initiateExport: " + e.message);
  }
}

/**
 * Polls for completion. Called by a time-based trigger.
 */
function pollExportStatus() {
  const startTime = new Date().getTime();
  const props = PropertiesService.getScriptProperties();
  const jobsJson = props.getProperty("EXPORT_JOBS");
  if (!jobsJson) {
    Logger.log("No active jobs found in properties. Stopping polling.");
    cleanupTriggers();
    return;
  }

  let jobs = JSON.parse(jobsJson);
  let pollCount = parseInt(props.getProperty("POLL_COUNT") || "0");
  const TOKEN = ScriptApp.getOAuthToken();
  const mboxFolder = DriveApp.getFolderById(CONFIG.MBOX_FOLDER_ID);

  Logger.log(`Running poll check #${pollCount + 1}...`);

  let anyChanged = false;
  let allFinished = true;

  for (let i = 0; i < jobs.length; i++) {
    let job = jobs[i];

    if (job.status === "IN_PROGRESS") {
      const statusUrl = `https://vault.googleapis.com/v1/matters/${CONFIG.MATTER_ID}/exports/${job.exportId}`;
      const statusResponse = UrlFetchApp.fetch(statusUrl, {
        method: "get",
        headers: { Authorization: `Bearer ${TOKEN}` },
        muteHttpExceptions: true
      });

      const completedExportData = JSON.parse(statusResponse.getContentText());
      job.status = completedExportData.status;
      Logger.log(`${job.type} Export Status: ${job.status}`);

      if (job.status === "COMPLETED") {
        anyChanged = true;
        processCompletedJob(job, completedExportData, mboxFolder, TOKEN);
        // If it switched to TRANSFERRING, we might want to start it now
        if (job.status === "TRANSFERRING") {
          allFinished = false;
          if (!isNearingTimeout(startTime)) {
            const finished = continueTransfer(job, TOKEN, startTime);
            if (finished) job.status = "COMPLETED_TRANSFERRED";
          }
        }
      } else if (job.status === "FAILED") {
        anyChanged = true;
        Logger.log(`ERROR: ${job.type} export failed.`);
      } else {
        allFinished = false;
      }
    } else if (job.status === "TRANSFERRING") {
      allFinished = false;
      Logger.log(`Resuming transfer for ${job.type} export...`);
      const finished = continueTransfer(job, TOKEN, startTime);
      anyChanged = true;
      if (finished) {
        job.status = "COMPLETED_TRANSFERRED";
      }
      if (isNearingTimeout(startTime)) break;
    } else if (job.status === "PENDING") {
      allFinished = false;
    } else if (job.status !== "COMPLETED_TRANSFERRED" && job.status !== "FAILED" && job.status !== "COMPLETED") {
      // Any other unexpected state
      allFinished = false;
    }
  }

  if (anyChanged) {
    props.setProperty("EXPORT_JOBS", JSON.stringify(jobs));
  }

  if (allFinished) {
    Logger.log("All jobs finished. Cleaning up triggers.");
    props.deleteProperty("EXPORT_JOBS");
    props.deleteProperty("POLL_COUNT");
    cleanupTriggers();
  } else {
    props.setProperty("POLL_COUNT", (pollCount + 1).toString());
    // Ensure trigger remains for next check
    setupTrigger();
  }
}

/**
 * Helper to process a single completed job.
 */
function processCompletedJob(job, completedExportData, mboxFolder, TOKEN) {
  if (completedExportData.cloudStorageSink && completedExportData.cloudStorageSink.files) {
    const files = completedExportData.cloudStorageSink.files;
    Logger.log(`Processing files for ${job.type} export...`);

    files.forEach(file => {
      const bucket = file.bucketName;
      const objectName = file.objectName;
      const fileName = objectName.split('/').pop();

      if (!fileName.toLowerCase().endsWith('.zip')) return;

      Logger.log(`Checking file size for ${job.type}: ${fileName}`);
      const downloadUrl = `https://storage.googleapis.com/storage/v1/b/${bucket}/o/${encodeURIComponent(objectName)}?alt=media`;
      
      const totalSize = getFileSize(downloadUrl, TOKEN);
      const sizeInMB = (totalSize / (1024 * 1024)).toFixed(2);
      Logger.log(`Total file size for ${job.type}: ${totalSize} bytes (${sizeInMB} MB)`);

      const targetMboxName = `${CONFIG.TARGET_USER}_${job.type}.mbox`;
      const targetZipName = `${CONFIG.TARGET_USER}_${job.type}_raw_export.zip`;

      if (totalSize >= 50 * 1024 * 1024) {
        Logger.log(`${job.type} export is large (${sizeInMB} MB). Setting up chunked transfer.`);
        const existingZips = mboxFolder.getFilesByName(targetZipName);
        while (existingZips.hasNext()) existingZips.next().setTrashed(true);
        
        const uploadUrl = initiateResumableUpload(targetZipName, mboxFolder.getId(), TOKEN);
        job.status = "TRANSFERRING";
        job.transferState = {
          downloadUrl: downloadUrl,
          uploadUrl: uploadUrl,
          start: 0,
          totalSize: totalSize,
          targetFileName: targetZipName
        };
        return;
      }

      const downloadResp = UrlFetchApp.fetch(downloadUrl, {
        method: "get",
        headers: { Authorization: `Bearer ${TOKEN}` },
        muteHttpExceptions: true
      });

      if (downloadResp.getResponseCode() === 200) {
        const zipBlob = downloadResp.getBlob();
        try {
          const unzippedBlobs = Utilities.unzip(zipBlob);
          let mboxSaved = false;
          unzippedBlobs.forEach(innerBlob => {
            if (innerBlob.getName().toLowerCase().endsWith(".mbox")) {
              const existingFiles = mboxFolder.getFilesByName(targetMboxName);
              while (existingFiles.hasNext()) existingFiles.next().setTrashed(true);
              const savedFile = mboxFolder.createFile(innerBlob);
              savedFile.setName(targetMboxName);
              Logger.log(`SUCCESS: Saved ${targetMboxName} to Drive.`);
              mboxSaved = true;
            }
          });
          if (!mboxSaved) Logger.log(`WARNING: No MBOX found in ${job.type} ZIP.`);
        } catch (zipErr) {
          Logger.log(`ERROR unzipping ${job.type}: ${zipErr.message}. Falling back to ZIP.`);
          const existingZips = mboxFolder.getFilesByName(targetZipName);
          while (existingZips.hasNext()) existingZips.next().setTrashed(true);
          mboxFolder.createFile(zipBlob).setName(targetZipName);
        }
      }
    });
  }
}

/**
 * Manages the poll trigger to avoid duplicates.
 */
function setupTrigger() {
  cleanupTriggers();
  ScriptApp.newTrigger("pollExportStatus")
    .timeBased()
    .after(2 * 60 * 1000) // 2 minutes
    .create();
}

function cleanupTriggers() {
  const triggers = ScriptApp.getProjectTriggers();
  triggers.forEach(t => {
    if (t.getHandlerFunction() === "pollExportStatus") {
      ScriptApp.deleteTrigger(t);
    }
  });
}

function exportMailboxAndSave() {
  initiateExport();
}

/**
 * Gets the total file size from GCS using a Range: bytes=0-0 request.
 */
function getFileSize(url, token) {
  const resp = UrlFetchApp.fetch(url, {
    method: "get",
    headers: { 
      Authorization: `Bearer ${token}`,
      Range: "bytes=0-0" 
    },
    muteHttpExceptions: true
  });

  const contentRange = resp.getHeaders()["Content-Range"] || resp.getHeaders()["content-range"];
  if (contentRange) {
    const size = contentRange.split("/")[1];
    return parseInt(size);
  }
  
  // Fallback to Content-Length if Range wasn't supported as expected
  const contentLength = resp.getHeaders()["Content-Length"] || resp.getHeaders()["content-length"];
  if (contentLength) return parseInt(contentLength);
  
  return 0;
}

/**
 * Initiates a resumable upload to Google Drive and returns the upload URL.
 */
function initiateResumableUpload(fileName, folderId, token) {
  const metadata = {
    name: fileName,
    parents: [folderId]
  };

  const response = UrlFetchApp.fetch(`https://www.googleapis.com/upload/drive/v3/files?uploadType=resumable&supportsAllDrives=true`, {
    method: "post",
    contentType: "application/json",
    headers: { Authorization: `Bearer ${token}` },
    payload: JSON.stringify(metadata),
    muteHttpExceptions: true
  });

  if (response.getResponseCode() !== 200) {
    throw new Error(`Failed to initiate resumable upload: ${response.getContentText()}`);
  }

  return response.getHeaders()["Location"] || response.getHeaders()["location"];
}

/**
 * Continues a chunked transfer until completion or timeout.
 */
function continueTransfer(job, token, startTime) {
  const state = job.transferState;
  const CHUNK_SIZE = 5 * 1024 * 1024; // Reduced to 5MB to avoid bandwidth quota issues

  while (state.start < state.totalSize) {
    if (isNearingTimeout(startTime)) {
      Logger.log(`Timeout approaching. Saving progress at ${state.start} bytes for ${state.targetFileName}`);
      return false; // Not finished
    }

    const end = Math.min(state.start + CHUNK_SIZE - 1, state.totalSize - 1);
    const range = `bytes=${state.start}-${end}`;

    Logger.log(`Streaming chunk ${range} for ${state.targetFileName}...`);

    const chunkResp = UrlFetchApp.fetch(state.downloadUrl, {
      method: "get",
      headers: { 
        Authorization: `Bearer ${token}`,
        Range: range 
      },
      muteHttpExceptions: true
    });

    if (chunkResp.getResponseCode() !== 206 && chunkResp.getResponseCode() !== 200) {
      throw new Error(`Failed to download chunk ${range}: ${chunkResp.getContentText()}`);
    }

    const chunkBytes = chunkResp.getContent();
    const uploadResp = UrlFetchApp.fetch(state.uploadUrl, {
      method: "put",
      headers: {
        Authorization: `Bearer ${token}`,
        "Content-Range": `bytes ${state.start}-${end}/${state.totalSize}`
      },
      payload: chunkBytes,
      muteHttpExceptions: true
    });

    if (uploadResp.getResponseCode() !== 200 && uploadResp.getResponseCode() !== 201 && uploadResp.getResponseCode() !== 308) {
      throw new Error(`Failed to upload chunk ${range}: ${uploadResp.getContentText()}`);
    }

    state.start += (end - state.start + 1); // Increment by actual bytes uploaded
    
    // Small delay between chunks to avoid "Bandwidth quota exceeded" errors
    Utilities.sleep(1000); 
  }

  Logger.log(`Successfully completed transfer of ${state.targetFileName}`);
  return true; // Finished
}

/**
 * Checks if the script is nearing its execution time limit.
 */
function isNearingTimeout(startTime) {
  const elapsed = new Date().getTime() - startTime;
  // Workspace limit is 30 mins.
  const limit = 30 * 60 * 1000;
  return elapsed > (limit - 60 * 1000); // 1 minute buffer
}