/**
 * Keys and tokens a task must never ship: the scrub that swaps them out of a task's
 * transcripts and run output, and the check that finds them anywhere else in the task.
 */

import { createHash } from 'crypto';
import { existsSync, readFileSync, readdirSync, statSync, writeFileSync } from 'fs';
import { join, relative } from 'path';

import { INPUT_CHECKSUMS_FILENAME } from './input-checksums';

export const SECRET_PLACEHOLDER = '[redacted]';

const SECRET_NAME_RE = /KEY|TOKEN|SECRET|PASSWORD/i;
// Short values are left alone so a scrub can't eat ordinary words.
const TOKEN_VALUE_RE = /^[A-Za-z0-9._~+/=-]{16,}$/;
const WORKER_KEY_RE = /sk-plwkr01-[A-Za-z0-9_-]{16,}/g;
// Full-length only: short `sk-ant-test-…` strings are fixtures that tasks are built around.
const DIRECT_KEY_RE = /sk-ant-[a-z]+\d{2}-[A-Za-z0-9_-]{80,}/g;
// Older worker keys are 24 uppercase letters and digits, so they're only recognizable after a key name.
const OLDER_KEY_RE =
  /(?<=(?:ANTHROPIC_API_KEY|OPENAI_API_KEY|[Xx]-[Aa][Pp][Ii]-[Kk][Ee][Yy])[^A-Za-z0-9]{1,6})[A-Z0-9]{24}(?![A-Za-z0-9_-])/g;
const MAX_SCAN_BYTES = 20 * 1024 * 1024;
// Files the worker didn't write: run output and detector reports.
const GENERATED_RE = /^(?:(?:reference-runs|rubric-regrades)\/[^/]+\/agent-output\/|detectors\/)/;
const TRANSCRIPT_RE = /^(?:session-full\.jsonl$|environment\/session|reference-runs\/|rubric-regrades\/)/;
const RESUMED_SESSION = 'environment/session.jsonl';

/** Credential values from the toolkit's `.env` and `explore/.env`. */
export function readEnvSecrets(root: string): string[] {
  const values = new Set<string>();
  for (const file of [join(root, '.env'), join(root, 'explore', '.env')]) {
    if (!existsSync(file)) continue;
    for (const line of readFileSync(file, 'utf8').split(/\r?\n/)) {
      const m = line.trim().match(/^(?:export\s+)?([A-Za-z_][A-Za-z0-9_]*)\s*=\s*(.*)$/);
      if (!m || !SECRET_NAME_RE.test(m[1])) continue;
      const value = m[2].trim().replace(/^(['"])(.*)\1$/, '$2');
      if (TOKEN_VALUE_RE.test(value)) values.add(value);
    }
  }
  return [...values];
}

function scrub(text: string, patterns: readonly RegExp[]): { text: string; count: number } {
  let count = 0;
  let out = text;
  for (const re of patterns) {
    out = out.replace(re, () => {
      count++;
      return SECRET_PLACEHOLDER;
    });
  }
  return { text: out, count };
}

const sha256 = (bytes: Buffer): string => createHash('sha256').update(bytes).digest('hex');

/** Restamp runs recorded against the unredacted resumed session, so a redaction alone never stales them. */
function restampRuns(taskDir: string, before: string, after: string): void {
  const runsDir = join(taskDir, 'reference-runs');
  if (!existsSync(runsDir)) return;
  for (const run of readdirSync(runsDir, { withFileTypes: true })) {
    const file = join(runsDir, run.name, INPUT_CHECKSUMS_FILENAME);
    if (!run.isDirectory() || !existsSync(file)) continue;
    try {
      const record = JSON.parse(readFileSync(file, 'utf8'));
      if (record?.inputs?.sessionJsonl !== before) continue;
      record.inputs.sessionJsonl = after;
      writeFileSync(file, JSON.stringify(record, null, 2) + '\n');
    } catch {
      // A malformed record already reads as unverifiable; leave it.
    }
  }
}

export interface TaskSecretScan {
  /** Transcripts, run output and detector reports whose keys were replaced in place. */
  scrubbed: string[];
  /** Files the worker wrote that hold a key, left for them to fix. */
  flagged: string[];
}

/** Scrub keys out of a task's transcripts and generated files in place, and list any other file that holds one. */
export function scrubTaskSecrets(taskDir: string, envSecrets: readonly string[]): TaskSecretScan {
  const exact = envSecrets.map((v) => new RegExp(v.replace(/[.*+?^${}()|[\]\\/]/g, '\\$&'), 'g'));
  // Code and fixtures can hold realistic fake keys of these shapes, so only transcripts are scrubbed for them.
  const strict = [...exact, WORKER_KEY_RE];
  const broad = [...strict, DIRECT_KEY_RE, OLDER_KEY_RE];
  const scan: TaskSecretScan = { scrubbed: [], flagged: [] };
  let session = null as { before: string; after: string } | null;

  const visit = (abs: string, rel: string): void => {
    const generated = GENERATED_RE.test(rel);
    const transcript = !generated && TRANSCRIPT_RE.test(rel);
    if (!transcript && statSync(abs).size > MAX_SCAN_BYTES) return;
    const bytes = readFileSync(abs);
    // latin1 maps each byte to one character, so every byte but the key's is written back as it was.
    const text = bytes.toString('latin1');
    if (generated || transcript) {
      const result = scrub(text, transcript ? broad : strict);
      if (result.count === 0) return;
      const scrubbed = Buffer.from(result.text, 'latin1');
      writeFileSync(abs, scrubbed);
      scan.scrubbed.push(rel);
      if (rel === RESUMED_SESSION) session = { before: sha256(bytes), after: sha256(scrubbed) };
    } else if (strict.some((re) => new RegExp(re.source).test(text))) {
      scan.flagged.push(rel);
    }
  };

  const walk = (dir: string): void => {
    for (const entry of readdirSync(dir, { withFileTypes: true })) {
      const abs = join(dir, entry.name);
      const rel = relative(taskDir, abs).split('\\').join('/');
      if (rel === 'environment/corpus' || entry.isSymbolicLink()) continue;
      if (entry.name === 'node_modules' || entry.name === '.git') continue;
      try {
        if (entry.isDirectory()) walk(abs);
        else if (entry.isFile()) visit(abs, rel);
      } catch {
        // Skipped: ingest runs this scrub again on the unpacked submission.
      }
    }
  };
  if (existsSync(taskDir)) walk(taskDir);
  if (session) restampRuns(taskDir, session.before, session.after);
  return scan;
}