#!/usr/bin/env bun

/**
 * Combined PDF Recovery Script
 *
 * Takes a corrupted base64-encoded PDF and produces the best recovered PDF.
 * Combines: text substitutions + stream-by-stream l→1 correction search
 *
 * Usage: bun recover_pdf.ts <input.txt> <output.pdf>
 */

import { readFileSync, writeFileSync, existsSync } from "fs";
import { inflateSync } from "zlib";
import { spawnSync } from "child_process";

// Check if a command exists
function commandExists(cmd: string): boolean {
  const result = spawnSync("which", [cmd], { encoding: "utf-8" });
  return result.status === 0;
}

// Run Ghostscript to re-distill PDF (makes it viewable)
function runGhostscript(inputPath: string, outputPath: string): boolean {
  if (!commandExists("gs")) {
    console.error("\n⚠ Ghostscript (gs) not found. Install with: brew install ghostscript");
    return false;
  }

  const result = spawnSync("gs", [
    "-sDEVICE=pdfwrite",
    "-dCompatibilityLevel=1.4",
    "-dNOPAUSE",
    "-dQUIET",
    "-dBATCH",
    `-sOutputFile=${outputPath}`,
    inputPath
  ], { encoding: "utf-8", stdio: ["pipe", "pipe", "pipe"] });

  return existsSync(outputPath);
}

// Extract text from PDF using Ghostscript
function extractText(pdfPath: string): string {
  if (!commandExists("gs")) {
    return "(gs not available for text extraction)";
  }

  const result = spawnSync("gs", [
    "-sDEVICE=txtwrite",
    "-sOutputFile=-",
    "-dQUIET",
    "-dNOPAUSE",
    "-dBATCH",
    pdfPath
  ], { encoding: "utf-8", stdio: ["pipe", "pipe", "pipe"] });

  return result.stdout || "";
}

// Known OCR error corrections (base64 level)
const BASE64_SUBSTITUTIONS: [string, string][] = [
  ["eGlw", "eG1w"],  // xipG → xmpG
  ["TmFp", "TmFt"],  // Naie → Name
  ["aW9k", "bW9k"],  // iode → mode
  ["Q0lZ", "Q01Z"],  // CIYK → CMYK
  ["UHJ1", "UHJl"],  // Pru → Pre (Predictor)
  ["b2JT", "b2NT"],  // ProbSet → ProcSet
  ["dElv", "dE1v"],  // LastIodified → LastModified
];

// Decoded content fixes (applied after base64 decode)
const DECODED_FIXES: [RegExp, string][] = [
  [/Naie\b/g, "Name"],
  [/\biode\b/g, "mode"],
  [/ProbSet/g, "ProcSet"],
  [/xipG/g, "xmpG"],
  [/>>\./g, ">>/"],
  [/\]\./g, "]/"],
  [/23\./g, "23/"],
  [/\.Type\b/g, "/Type"],
  [/\.Prop/g, "/Prop"],
  [/\.Reso/g, "/Reso"],
  [/\.Filter\b/g, "/Filter"],
  [/\.Length\b/g, "/Length"],
  [/\.Pages\b/g, "/Pages"],
  [/\.Root\b/g, "/Root"],
  [/\.Info\b/g, "/Info"],
  [/\.Size\b/g, "/Size"],
  [/%4338D8C/g, "54338D8C"],  // Hex ID fix
];

const L_BYTE = 'l'.charCodeAt(0);
const ONE_BYTE = '1'.charCodeAt(0);

interface StreamLoc {
  obj: number;
  start: number;
  end: number;
  size: number;
}

function pdfToB64(pdfPos: number): number {
  return Math.floor((pdfPos * 4) / 3);
}

function findObjStream(pdfStr: string, objNum: number): StreamLoc | null {
  const pattern = new RegExp(objNum + '\\s+0\\s+obj[^>]*>>\\s*stream[\\r\\n]+');
  const match = pdfStr.match(pattern);
  if (!match) return null;

  const start = match.index! + match[0].length;
  const end = pdfStr.indexOf('endstream', start);
  if (end === -1) return null;
  return { obj: objNum, start, end, size: end - start };
}

function canDecompress(pdf: Buffer, stream: StreamLoc): boolean {
  try {
    inflateSync(pdf.subarray(stream.start, stream.end));
    return true;
  } catch {
    return false;
  }
}

function findAllStreams(pdfStr: string): StreamLoc[] {
  const streams: StreamLoc[] = [];
  const pattern = /(\d+)\s+0\s+obj[^>]*>>\s*stream[\r\n]+/g;
  let match;

  while ((match = pattern.exec(pdfStr)) !== null) {
    const objNum = parseInt(match[1]);
    const start = match.index + match[0].length;
    const end = pdfStr.indexOf('endstream', start);
    if (end !== -1) {
      streams.push({ obj: objNum, start, end, size: end - start });
    }
  }

  return streams;
}

async function main() {
  const args = process.argv.slice(2);

  if (args.length < 2) {
    console.log("Usage: bun recover_pdf.ts <input.txt> <output.pdf>");
    console.log("  input.txt  - base64-encoded PDF file");
    console.log("  output.pdf - recovered PDF output");
    process.exit(1);
  }

  const inputFile = args[0];
  const outputFile = args[1];

  console.log("╔══════════════════════════════════════════════════════════════╗");
  console.log("║       Combined PDF Recovery Script                           ║");
  console.log("╚══════════════════════════════════════════════════════════════╝");

  // Step 1: Read original base64
  console.log(`\n[1/4] Reading ${inputFile}...`);
  let base64Str = readFileSync(inputFile, "utf-8").trim();
  console.log(`  ${base64Str.length} characters`);

  // Step 2: Apply known text substitutions
  console.log("\n[2/4] Applying known text corrections...");
  for (const [from, to] of BASE64_SUBSTITUTIONS) {
    const count = (base64Str.match(new RegExp(from, 'g')) || []).length;
    if (count > 0) {
      base64Str = base64Str.replace(new RegExp(from, 'g'), to);
      console.log(`  ${from} → ${to}: ${count} fixes`);
    }
  }

  let base64Buf = Buffer.from(base64Str);
  let pdf = Buffer.from(base64Buf.toString(), 'base64');
  let pdfStr = pdf.toString('binary');

  // Apply decoded content fixes
  console.log("\n  Applying decoded content fixes...");
  for (const [pattern, replacement] of DECODED_FIXES) {
    const matches = pdfStr.match(pattern);
    if (matches && matches.length > 0) {
      pdfStr = pdfStr.replace(pattern, replacement);
      console.log(`    ${pattern.source}: ${matches.length} fixes`);
    }
  }
  pdf = Buffer.from(pdfStr, 'binary');
  base64Buf = Buffer.from(pdf.toString('base64'));

  // Step 3: Find all streams and check status
  console.log("\n[3/4] Analyzing streams...");
  const allStreams = findAllStreams(pdfStr);
  let okCount = allStreams.filter(s => canDecompress(pdf, s)).length;
  console.log(`  Found ${allStreams.length} streams, ${okCount} decompress OK`);

  // Page content streams (most important for visible content)
  const pageContentObjs = [39, 40, 41, 42, 43, 44, 45, 46];

  console.log("\n  Page content streams:");
  for (const objNum of pageContentObjs) {
    const s = findObjStream(pdfStr, objNum);
    if (s) {
      const ok = canDecompress(pdf, s);
      console.log(`    Object ${objNum}: ${s.size}b [${ok ? 'OK' : 'FAIL'}]`);
    }
  }

  // Step 4: Repair streams with l→1 corrections
  console.log("\n[4/4] Repairing streams...");

  // Only attempt streams that are feasible to brute-force (< 3000 bytes)
  const MAX_STREAM_SIZE = 3000;

  // Get all failed streams, prioritize page content, sort by size
  const failedStreams = allStreams
    .filter(s => !canDecompress(pdf, s) && s.size < MAX_STREAM_SIZE)
    .sort((a, b) => {
      // Page content first
      const aIsPage = pageContentObjs.includes(a.obj);
      const bIsPage = pageContentObjs.includes(b.obj);
      if (aIsPage && !bIsPage) return -1;
      if (!aIsPage && bIsPage) return 1;
      // Then by size (smaller first)
      return a.size - b.size;
    });

  console.log(`  ${failedStreams.length} streams to attempt (< ${MAX_STREAM_SIZE}b)\n`);

  let totalFixed = 0;

  for (const stream of failedStreams) {
    // Re-check in case earlier fixes affected this stream
    pdf = Buffer.from(base64Buf.toString(), 'base64');
    pdfStr = pdf.toString('binary');
    const currentStream = findObjStream(pdfStr, stream.obj);
    if (!currentStream || canDecompress(pdf, currentStream)) continue;

    const isPageContent = pageContentObjs.includes(stream.obj);
    const label = isPageContent ? '[PAGE]' : '[other]';
    process.stdout.write(`  Object ${stream.obj} ${label} (${currentStream.size}b): `);

    const padding = 200;
    const b64Start = Math.max(0, pdfToB64(currentStream.start) - padding);
    const b64End = Math.min(base64Buf.length, pdfToB64(currentStream.end) + padding);

    // Find all 'l' positions in this region
    const lPos: number[] = [];
    for (let i = b64Start; i < b64End; i++) {
      if (base64Buf[i] === L_BYTE) lPos.push(i);
    }

    let found = false;

    // Try single l→1
    for (const p of lPos) {
      base64Buf[p] = ONE_BYTE;
      pdf = Buffer.from(base64Buf.toString(), 'base64');
      const s = findObjStream(pdf.toString('binary'), stream.obj);
      if (s && canDecompress(pdf, s)) {
        console.log(`1x l→1`);
        found = true;
        totalFixed++;
        break;
      }
      base64Buf[p] = L_BYTE;
    }
    if (found) continue;

    // Try double l→1
    outer2: for (let i = 0; i < lPos.length; i++) {
      for (let j = i + 1; j < lPos.length; j++) {
        base64Buf[lPos[i]] = ONE_BYTE;
        base64Buf[lPos[j]] = ONE_BYTE;
        pdf = Buffer.from(base64Buf.toString(), 'base64');
        const s = findObjStream(pdf.toString('binary'), stream.obj);
        if (s && canDecompress(pdf, s)) {
          console.log(`2x l→1`);
          found = true;
          totalFixed++;
          break outer2;
        }
        base64Buf[lPos[i]] = L_BYTE;
        base64Buf[lPos[j]] = L_BYTE;
      }
    }
    if (found) continue;

    // Try triple l→1 (limit to 50 positions)
    if (lPos.length <= 50) {
      outer3: for (let i = 0; i < lPos.length; i++) {
        for (let j = i + 1; j < lPos.length; j++) {
          for (let k = j + 1; k < lPos.length; k++) {
            base64Buf[lPos[i]] = ONE_BYTE;
            base64Buf[lPos[j]] = ONE_BYTE;
            base64Buf[lPos[k]] = ONE_BYTE;
            pdf = Buffer.from(base64Buf.toString(), 'base64');
            const s = findObjStream(pdf.toString('binary'), stream.obj);
            if (s && canDecompress(pdf, s)) {
              console.log(`3x l→1`);
              found = true;
              totalFixed++;
              break outer3;
            }
            base64Buf[lPos[i]] = L_BYTE;
            base64Buf[lPos[j]] = L_BYTE;
            base64Buf[lPos[k]] = L_BYTE;
          }
        }
      }
    }
    if (found) continue;

    // Try quad l→1 (limit to 40 positions)
    if (lPos.length <= 40) {
      outer4: for (let i = 0; i < lPos.length; i++) {
        for (let j = i + 1; j < lPos.length; j++) {
          for (let k = j + 1; k < lPos.length; k++) {
            for (let m = k + 1; m < lPos.length; m++) {
              base64Buf[lPos[i]] = ONE_BYTE;
              base64Buf[lPos[j]] = ONE_BYTE;
              base64Buf[lPos[k]] = ONE_BYTE;
              base64Buf[lPos[m]] = ONE_BYTE;
              pdf = Buffer.from(base64Buf.toString(), 'base64');
              const s = findObjStream(pdf.toString('binary'), stream.obj);
              if (s && canDecompress(pdf, s)) {
                console.log(`4x l→1`);
                found = true;
                totalFixed++;
                break outer4;
              }
              base64Buf[lPos[i]] = L_BYTE;
              base64Buf[lPos[j]] = L_BYTE;
              base64Buf[lPos[k]] = L_BYTE;
              base64Buf[lPos[m]] = L_BYTE;
            }
          }
        }
      }
    }

    if (!found) {
      console.log(`no fix (${lPos.length} 'l' positions)`);
    }
  }

  // Final output
  pdf = Buffer.from(base64Buf.toString(), 'base64');
  pdfStr = pdf.toString('binary');

  console.log("\n╔══════════════════════════════════════════════════════════════╗");
  console.log("║                       RESULTS                                ║");
  console.log("╚══════════════════════════════════════════════════════════════╝");

  const finalStreams = findAllStreams(pdfStr);
  okCount = finalStreams.filter(s => canDecompress(pdf, s)).length;
  console.log(`\nTotal streams: ${okCount}/${finalStreams.length} OK`);
  console.log(`Streams fixed: ${totalFixed}`);

  // Page content status
  let pageOk = 0;
  console.log("\nPage content streams:");
  for (const objNum of pageContentObjs) {
    const s = findObjStream(pdfStr, objNum);
    if (s) {
      const ok = canDecompress(pdf, s);
      if (ok) pageOk++;
      console.log(`  Object ${objNum}: ${ok ? '✓' : '✗'}`);
    }
  }
  console.log(`\nPage content: ${pageOk}/${pageContentObjs.length} OK`);

  // Save raw repaired PDF
  const rawFile = outputFile.replace(/\.pdf$/, "-raw.pdf");
  writeFileSync(rawFile, pdf);
  console.log(`\nSaved raw: ${rawFile}`);

  // Re-distill with Ghostscript to make it viewable
  console.log("\n[5/5] Re-distilling with Ghostscript...");
  if (runGhostscript(rawFile, outputFile)) {
    console.log(`  ✓ Saved: ${outputFile}`);

    // Extract and show text
    console.log("\n--- Extracted Text ---");
    const text = extractText(outputFile);
    const lines = text.split('\n').filter(l => l.trim()).slice(0, 25);
    if (lines.length > 0) {
      console.log(lines.join('\n'));
    } else {
      console.log("(no text extracted)");
    }
  } else {
    console.log(`  ✗ Ghostscript re-distill failed`);
    console.log(`  Raw PDF saved to: ${rawFile}`);
    console.log("\n  To install Ghostscript: brew install ghostscript");
  }
}

main().catch(console.error);