import { readdir, readFile } from "node:fs/promises";
import { join } from "node:path";
import { canonicalize } from "../../schema/src/canonical.js";
import { atomicWriteBytes } from "./artifacts.js";
import type {
  CellEvidence,
  CellEvidencePayload,
  ConsoleEntry,
  NetworkEntry,
} from "./capture.js";

export type RedactionRule = {
  name: string;
  secret: string;
  replacement?: string;
};

export type RedactionReport = {
  rules: Array<{
    name: string;
    replacements: number;
  }>;
};

type ReplacementCounts = Map<string, number>;

function assertNonEmptyString(value: string, field: string): void {
  if (value.trim().length === 0) {
    throw new Error(`${field} must be a non-empty string`);
  }
}

export function validateRedactionRules(rules: RedactionRule[]): void {
  for (const rule of rules) {
    assertNonEmptyString(rule.name, "redaction rule name");
    assertNonEmptyString(rule.secret, `redaction rule ${rule.name} secret`);
    if (rule.replacement !== undefined) {
      assertNonEmptyString(rule.replacement, `redaction rule ${rule.name} replacement`);
    }
  }
}

function initReplacementCounts(rules: RedactionRule[]): ReplacementCounts {
  const counts = new Map<string, number>();
  for (const rule of rules) {
    counts.set(rule.name, 0);
  }
  return counts;
}

function countOccurrences(text: string, secret: string): number {
  if (secret.length === 0) {
    return 0;
  }
  let count = 0;
  let index = text.indexOf(secret);
  while (index !== -1) {
    count += 1;
    index = text.indexOf(secret, index + secret.length);
  }
  return count;
}

function redactText(
  text: string,
  rules: RedactionRule[],
  counts: ReplacementCounts,
): string {
  let redacted = text;
  for (const rule of rules) {
    const token = rule.replacement ?? `[redacted:${rule.name}]`;
    const matches = countOccurrences(redacted, rule.secret);
    if (matches > 0) {
      counts.set(rule.name, (counts.get(rule.name) ?? 0) + matches);
      redacted = redacted.split(rule.secret).join(token);
    }
  }
  return redacted;
}

function redactHeaderRecord(
  headers: Record<string, string> | undefined,
  rules: RedactionRule[],
  counts: ReplacementCounts,
): Record<string, string> | undefined {
  if (headers === undefined) {
    return undefined;
  }
  const redacted: Record<string, string> = {};
  for (const [key, value] of Object.entries(headers)) {
    redacted[key] = redactText(value, rules, counts);
  }
  return redacted;
}

function redactConsoleEntry(
  entry: ConsoleEntry,
  rules: RedactionRule[],
  counts: ReplacementCounts,
): ConsoleEntry {
  const redacted: ConsoleEntry = {
    type: entry.type,
    text: redactText(entry.text, rules, counts),
  };
  if (entry.stack !== undefined) {
    redacted.stack = redactText(entry.stack, rules, counts);
  }
  if (entry.location !== undefined) {
    const location = { ...entry.location };
    if (location.url !== undefined) {
      location.url = redactText(location.url, rules, counts);
    }
    redacted.location = location;
  }
  return redacted;
}

function redactNetworkEntry(
  entry: NetworkEntry,
  rules: RedactionRule[],
  counts: ReplacementCounts,
): NetworkEntry {
  const redacted: NetworkEntry = {
    url: redactText(entry.url, rules, counts),
    method: entry.method,
    resourceType: entry.resourceType,
  };
  if (entry.status !== undefined) {
    redacted.status = entry.status;
  }
  const requestHeaders = redactHeaderRecord(entry.requestHeaders, rules, counts);
  if (requestHeaders !== undefined) {
    redacted.requestHeaders = requestHeaders;
  }
  const responseHeaders = redactHeaderRecord(entry.responseHeaders, rules, counts);
  if (responseHeaders !== undefined) {
    redacted.responseHeaders = responseHeaders;
  }
  if (entry.requestBody !== undefined) {
    redacted.requestBody = redactText(entry.requestBody, rules, counts);
  }
  if (entry.responseBody !== undefined) {
    redacted.responseBody = redactText(entry.responseBody, rules, counts);
  }
  return redacted;
}

export function redactEvidencePayload(
  payload: CellEvidencePayload,
  rules: RedactionRule[],
): {
  payload: CellEvidencePayload;
  counts: ReplacementCounts;
  artifactReplacements: { domSnapshot: number; trace: number };
} {
  validateRedactionRules(rules);
  const counts = initReplacementCounts(rules);

  const domBefore = payload.domSnapshot.bytes.toString("utf8");
  const domAfter = redactText(domBefore, rules, counts);
  const domReplacements = rules.reduce(
    (total, rule) => total + countOccurrences(domBefore, rule.secret),
    0,
  );

  let traceReplacements = 0;
  let tracePayload = payload.trace;
  if (payload.trace !== undefined) {
    const traceBefore = payload.trace.bytes.toString("utf8");
    const traceAfter = redactText(traceBefore, rules, counts);
    traceReplacements = rules.reduce(
      (total, rule) => total + countOccurrences(traceBefore, rule.secret),
      0,
    );
    tracePayload = {
      ...payload.trace,
      bytes: Buffer.from(traceAfter, "utf8"),
    };
  }

  const redactedPayload: CellEvidencePayload = {
    ...payload,
    console: payload.console.map((entry) => redactConsoleEntry(entry, rules, counts)),
    network: payload.network.map((entry) => redactNetworkEntry(entry, rules, counts)),
    domSnapshot: {
      ...payload.domSnapshot,
      bytes: Buffer.from(domAfter, "utf8"),
    },
    ...(tracePayload === undefined ? {} : { trace: tracePayload }),
  };

  return {
    payload: redactedPayload,
    counts,
    artifactReplacements: {
      domSnapshot: domReplacements,
      trace: traceReplacements,
    },
  };
}

export function redactJsonPayload(
  value: unknown,
  rules: RedactionRule[],
): { json: string; replacements: number } {
  validateRedactionRules(rules);
  const counts = initReplacementCounts(rules);
  const serialized = canonicalize(value);
  const json = redactText(serialized, rules, counts);
  const replacements = [...counts.values()].reduce((total, count) => total + count, 0);
  return { json, replacements };
}

export function redactEvidence(
  evidence: CellEvidence,
  rules: RedactionRule[],
): CellEvidence {
  validateRedactionRules(rules);
  const counts = initReplacementCounts(rules);

  return {
    ...evidence,
    console: evidence.console.map((entry) => redactConsoleEntry(entry, rules, counts)),
    network: evidence.network.map((entry) => redactNetworkEntry(entry, rules, counts)),
  };
}

function buildReport(counts: ReplacementCounts, rules: RedactionRule[]): RedactionReport {
  return {
    rules: rules.map((rule) => ({
      name: rule.name,
      replacements: counts.get(rule.name) ?? 0,
    })),
  };
}

function isImageArtifactPath(relativePath: string): boolean {
  return /\.(?:gif|jpe?g|png|svg|webp)$/i.test(relativePath);
}

async function walkFiles(rootDir: string): Promise<string[]> {
  const files: string[] = [];

  async function visit(currentDir: string): Promise<void> {
    const entries = await readdir(currentDir, { withFileTypes: true });
    for (const entry of entries) {
      const absolutePath = join(currentDir, entry.name);
      if (entry.isDirectory()) {
        await visit(absolutePath);
        continue;
      }
      if (entry.isFile()) {
        files.push(absolutePath);
      }
    }
  }

  await visit(rootDir);
  return files;
}

export async function redactArtifacts(
  runDir: string,
  rules: RedactionRule[],
): Promise<RedactionReport> {
  assertNonEmptyString(runDir, "runDir");
  validateRedactionRules(rules);

  const counts = initReplacementCounts(rules);
  const files = await walkFiles(runDir);

  for (const filePath of files) {
    if (isImageArtifactPath(filePath)) {
      continue;
    }

    const original = await readFile(filePath);
    const originalText = original.toString("utf8");
    const redactedText = redactText(originalText, rules, counts);
    if (redactedText !== originalText) {
      await atomicWriteBytes(filePath, Buffer.from(redactedText, "utf8"));
    }
  }

  return buildReport(counts, rules);
}

export function redactionStateForArtifact(replacements: number): "none" | "partial" | "full" {
  if (replacements === 0) {
    return "none";
  }
  return "partial";
}