#!/usr/bin/env node
'use strict';

const fs = require('fs');
const os = require('os');
const path = require('path');
const { spawnSync } = require('child_process');

const DEFAULT_MAX_LONG_EDGE = 2000;

function usage() {
  console.error('Usage: repair-claude-session-images.js <session-jsonl-or-session-id> [--max-long-edge=2000] [--dry-run]');
}

function parseArgs(argv) {
  const args = { dryRun: false, maxLongEdge: DEFAULT_MAX_LONG_EDGE, target: '' };
  for (const arg of argv) {
    if (arg === '--dry-run') args.dryRun = true;
    else if (arg.startsWith('--max-long-edge=')) args.maxLongEdge = Number(arg.slice('--max-long-edge='.length)) || DEFAULT_MAX_LONG_EDGE;
    else if (!args.target) args.target = arg;
    else throw new Error('Unexpected argument: ' + arg);
  }
  if (!args.target) throw new Error('Missing session path or id');
  return args;
}

function findSessionJsonl(target) {
  if (target.endsWith('.jsonl') || target.includes(path.sep)) {
    const resolved = path.resolve(target.replace(/^~/, os.homedir()));
    if (!fs.existsSync(resolved)) throw new Error('Session JSONL not found: ' + resolved);
    return resolved;
  }
  const root = path.join(os.homedir(), '.claude', 'projects');
  const stack = [root];
  const suffix = target + '.jsonl';
  while (stack.length) {
    const dir = stack.pop();
    let entries = [];
    try { entries = fs.readdirSync(dir, { withFileTypes: true }); } catch { continue; }
    for (const ent of entries) {
      const p = path.join(dir, ent.name);
      if (ent.isDirectory()) stack.push(p);
      else if (ent.isFile() && ent.name === suffix) return p;
    }
  }
  throw new Error('Session id not found under ~/.claude/projects: ' + target);
}

function imageDimensions(buffer) {
  if (!Buffer.isBuffer(buffer) || buffer.length < 24) return null;
  if (buffer.slice(0, 8).equals(Buffer.from([137, 80, 78, 71, 13, 10, 26, 10]))) {
    return { type: 'png', width: buffer.readUInt32BE(16), height: buffer.readUInt32BE(20) };
  }
  if (buffer[0] === 0xff && buffer[1] === 0xd8) {
    let offset = 2;
    while (offset + 8 < buffer.length) {
      while (buffer[offset] === 0xff) offset += 1;
      const marker = buffer[offset++];
      if (marker === 0xda || marker === 0xd9) break;
      const length = buffer.readUInt16BE(offset);
      if (marker >= 0xc0 && marker <= 0xcf && ![0xc4, 0xc8, 0xcc].includes(marker)) {
        return { type: 'jpeg', width: buffer.readUInt16BE(offset + 5), height: buffer.readUInt16BE(offset + 3) };
      }
      offset += length;
    }
  }
  return null;
}

function extensionForMediaType(mediaType, detectedType) {
  const m = String(mediaType || detectedType || '').toLowerCase();
  if (m.includes('jpeg') || m.includes('jpg')) return '.jpg';
  return '.png';
}

function runSipsResize(buffer, mediaType, maxLongEdge) {
  const detected = imageDimensions(buffer);
  const ext = extensionForMediaType(mediaType, detected && detected.type);
  const dir = fs.mkdtempSync(path.join(os.tmpdir(), 'ctm-image-repair-'));
  const input = path.join(dir, 'input' + ext);
  const output = path.join(dir, 'output' + ext);
  try {
    fs.writeFileSync(input, buffer);
    const result = spawnSync('sips', ['-Z', String(maxLongEdge), input, '--out', output], {
      encoding: 'utf8',
      maxBuffer: 1024 * 1024,
    });
    if (result.status !== 0) {
      throw new Error((result.stderr || result.stdout || 'sips failed').trim());
    }
    const out = fs.readFileSync(output);
    const dims = imageDimensions(out);
    if (!dims || dims.width > maxLongEdge || dims.height > maxLongEdge) {
      throw new Error('resize did not produce a valid capped image');
    }
    return { buffer: out, dimensions: dims };
  } finally {
    try { fs.rmSync(dir, { recursive: true, force: true }); } catch {}
  }
}

function visitImageBlocks(value, visitor) {
  if (!value || typeof value !== 'object') return;
  if (Array.isArray(value)) {
    value.forEach(item => visitImageBlocks(item, visitor));
    return;
  }
  if (value.type === 'image' && value.source && value.source.type === 'base64' && value.source.data) {
    visitor(value);
  }
  for (const key of Object.keys(value)) visitImageBlocks(value[key], visitor);
}

function scrubBase64ImageSource(source) {
  if (!source || source.type !== 'base64') return [];
  const allowed = new Set(['type', 'media_type', 'data']);
  const removed = [];
  for (const key of Object.keys(source)) {
    if (allowed.has(key)) continue;
    removed.push(key);
    delete source[key];
  }
  return removed;
}

function repairFile(file, opts) {
  const original = fs.readFileSync(file, 'utf8');
  const lines = original.split(/\n/);
  const repairs = [];
  const out = lines.map((line, lineIndex) => {
    if (!line.trim()) return line;
    let parsed;
    try { parsed = JSON.parse(line); } catch { return line; }
    let changed = false;
    visitImageBlocks(parsed, (block) => {
      const scrubbedKeys = scrubBase64ImageSource(block.source);
      if (scrubbedKeys.length) {
        repairs.push({
          line: lineIndex + 1,
          scrubbedKeys,
          action: 'scrub-invalid-image-source-metadata',
        });
        if (!opts.dryRun) changed = true;
      }
      let buffer;
      try { buffer = Buffer.from(block.source.data, 'base64'); } catch { return; }
      const before = imageDimensions(buffer);
      if (!before || Math.max(before.width, before.height) <= opts.maxLongEdge) return;
      repairs.push({
        line: lineIndex + 1,
        before,
        mediaType: block.source.media_type || '',
        bytesBefore: buffer.length,
      });
      if (opts.dryRun) return;
      const resized = runSipsResize(buffer, block.source.media_type, opts.maxLongEdge);
      block.source.data = resized.buffer.toString('base64');
      changed = true;
      repairs[repairs.length - 1].after = resized.dimensions;
      repairs[repairs.length - 1].bytesAfter = resized.buffer.length;
    });
    return changed ? JSON.stringify(parsed) : line;
  });
  if (repairs.length && !opts.dryRun) {
    const stamp = new Date().toISOString().replace(/[:.]/g, '-');
    const backup = file + '.ctm-image-backup-' + stamp;
    fs.copyFileSync(file, backup);
    const tmp = file + '.tmp-' + process.pid;
    fs.writeFileSync(tmp, out.join('\n'), 'utf8');
    fs.renameSync(tmp, file);
    return { file, backup, repairs };
  }
  return { file, backup: null, repairs };
}

function main() {
  try {
    const args = parseArgs(process.argv.slice(2));
    const file = findSessionJsonl(args.target);
    const result = repairFile(file, args);
    console.log(JSON.stringify(result, null, 2));
    if (!result.repairs.length) process.exitCode = 0;
  } catch (err) {
    usage();
    console.error(err && err.message ? err.message : err);
    process.exit(1);
  }
}

if (require.main === module) main();

module.exports = {
  imageDimensions,
  repairFile,
  scrubBase64ImageSource,
};
