import { createHash } from "node:crypto";
import type { Page } from "playwright";
import { capturePageScreenshot } from "./screenshot.js";
import {
  contentAddressedFileName,
  persistArtifactAtPath,
  type ArtifactRef,
} from "./artifacts.js";
import { validateRedactionRules, type RedactionRule } from "./redaction.js";

export type CellOriginalRef = {
  relativePath: string;
  contentHash: string;
  mediaType: string;
  dimensions?: { width: number; height: number };
};

export type CellOriginalCapture = {
  ref: CellOriginalRef;
  pageBounds: { width: number; height: number };
  sourceRunId?: string;
};

export type CaptureCellOriginalOptions = {
  fileName?: string;
  artifactRunDir?: string;
  sourceRunId?: string;
  redactionRules?: RedactionRule[];
};

function sha256Hex(bytes: Buffer): string {
  return createHash("sha256").update(bytes).digest("hex");
}

export async function captureCellOriginal(
  page: Page,
  cellId: string,
  options: CaptureCellOriginalOptions = {},
): Promise<CellOriginalCapture> {
  validateRedactionRules(options.redactionRules ?? []);
  const pageBounds = await page.evaluate(() => ({
    width: document.documentElement.scrollWidth,
    height: document.documentElement.scrollHeight,
  }));
  const bytes = await capturePageScreenshot(page);
  const viewport = page.viewportSize();
  const dimensions =
    pageBounds.width > 0 && pageBounds.height > 0
      ? { width: pageBounds.width, height: pageBounds.height }
      : viewport === null
        ? undefined
        : { width: viewport.width, height: viewport.height };
  const contentHash = sha256Hex(bytes);
  const ref: CellOriginalRef = {
    relativePath: `cells/${cellId}/${contentAddressedFileName(options.fileName ?? "screenshot.png", contentHash)}`,
    contentHash,
    mediaType: "image/png",
    ...(dimensions === undefined ? {} : { dimensions }),
  };

  if (options.artifactRunDir !== undefined && options.sourceRunId !== undefined) {
    await persistArtifactAtPath({
      runDir: options.artifactRunDir,
      relativePath: ref.relativePath,
      payload: {
        bytes,
        mediaType: ref.mediaType,
        ...(dimensions === undefined ? {} : { dimensions }),
      },
      sourceRunId: options.sourceRunId,
    });
    return { ref, pageBounds, sourceRunId: options.sourceRunId };
  }

  return { ref, pageBounds };
}

export type CellScreenshotInput = {
  cell: { id: string };
  capture: { screenshot: boolean };
  artifactRunDir?: string;
  sourceRunId?: string;
  redactionRules?: RedactionRule[];
};

export async function captureCellScreenshots(
  page: Page,
  input: CellScreenshotInput,
): Promise<{ original: ArtifactRef } | undefined> {
  if (!input.capture.screenshot) {
    return undefined;
  }
  const { artifactRunDir, sourceRunId } = input;
  if (artifactRunDir === undefined || sourceRunId === undefined) {
    return undefined;
  }
  const captured = await captureCellOriginal(page, input.cell.id, {
    artifactRunDir,
    sourceRunId,
    ...(input.redactionRules === undefined ? {} : { redactionRules: input.redactionRules }),
  });
  return { original: cellOriginalArtifactRef(captured.ref, sourceRunId) };
}

export function cellScreenshotArtifacts(evidence: {
  screenshots?: { original: ArtifactRef };
}): ArtifactRef[] {
  const original = evidence.screenshots?.original;
  return original === undefined ? [] : [original];
}

export function cellOriginalArtifactRef(ref: CellOriginalRef, sourceRunId: string): ArtifactRef {
  const artifact: ArtifactRef = {
    relativePath: ref.relativePath,
    mediaType: ref.mediaType,
    contentHash: ref.contentHash,
    redactionState: "none",
    sourceRunId,
  };
  if (ref.dimensions !== undefined) {
    artifact.dimensions = ref.dimensions;
  }
  return artifact;
}