import { assert, ErrorWithExtra, unreachable } from '../../../common/util/util.js';
import {
  EncodableTextureFormat,
  getTextureFormatType,
  isColorTextureFormat,
  isDepthTextureFormat,
} from '../../format_info.js';
import { GPUTestBase } from '../../gpu_test.js';
import { numbersApproximatelyEqual } from '../conversion.js';
import { generatePrettyTable, numericToStringBuilder } from '../pretty_diff_tables.js';
import { reifyExtent3D, reifyOrigin3D } from '../unions.js';

import { fullSubrectCoordinates } from './base.js';
import { getTextureSubCopyLayout } from './layout.js';
import { kTexelRepresentationInfo, PerTexelComponent, TexelComponent } from './texel_data.js';
import { TexelView } from './texel_view.js';

type PerPixelAtLevel<T> = (coords: Required<GPUOrigin3DDict>) => T;

/** Threshold options for comparing texels of different formats (norm/float/int). */
export type TexelCompareOptions = {
  /** Threshold for integer texture formats. Defaults to 0. */
  maxIntDiff?: number;
  /** Threshold for non-integer (norm/float) texture formats, if not overridden. */
  maxFractionalDiff?: number;
  /** Threshold in ULPs for unorm/snorm texture formats. Overrides `maxFractionalDiff`. */
  maxDiffULPsForNormFormat?: number;
  /** Threshold in ULPs for float/ufloat texture formats. Overrides `maxFractionalDiff`. */
  maxDiffULPsForFloatFormat?: number;
};

export type PixelExpectation = PerTexelComponent<number> | Uint8Array;

export type PerPixelComparison<E extends PixelExpectation> = {
  coord: GPUOrigin3D;
  exp: E;
};

type TexelViewComparer = {
  /** Given coords, returns whether the two texel views are considered matching at that point. */
  predicate: PerPixelAtLevel<boolean>;
  /**
   * Given a list of failed coords, returns table rows for `generatePrettyTable` that
   * display the actual/expected values and diffs for debugging.
   */
  tableRows: (failedCoords: readonly Required<GPUOrigin3DDict>[]) => Iterable<string>[];
};

// Creates a texel view comparer that compares an actual texel view against
// an expected texel view and optionally a second, alternate, texel view.
// The comparison for each texel succeeds if the actual texel is close enough
// to either
function makeTexelViewComparer(
  format: EncodableTextureFormat,
  {
    actTexelView,
    expTexelView,
    altExpTexelView,
  }: { actTexelView: TexelView; expTexelView: TexelView; altExpTexelView?: TexelView },
  opts: TexelCompareOptions
): TexelViewComparer {
  const {
    maxIntDiff = 0,
    maxFractionalDiff,
    maxDiffULPsForNormFormat,
    maxDiffULPsForFloatFormat,
  } = opts;

  assert(maxIntDiff >= 0, 'threshold must be non-negative');
  if (maxFractionalDiff !== undefined) {
    assert(maxFractionalDiff >= 0, 'threshold must be non-negative');
  }
  if (maxDiffULPsForFloatFormat !== undefined) {
    assert(maxDiffULPsForFloatFormat >= 0, 'threshold must be non-negative');
  }
  if (maxDiffULPsForNormFormat !== undefined) {
    assert(maxDiffULPsForNormFormat >= 0, 'threshold must be non-negative');
  }

  const fmtIsInt = format.includes('int');
  const fmtIsNorm = format.includes('norm');
  const fmtIsFloat = format.includes('float');

  const tvc = {} as TexelViewComparer;
  if (fmtIsInt) {
    tvc.predicate = coords =>
      comparePerComponent(
        actTexelView.color(coords),
        expTexelView.color(coords),
        maxIntDiff,
        altExpTexelView?.color(coords)
      );
  } else if (fmtIsNorm && maxDiffULPsForNormFormat !== undefined) {
    tvc.predicate = coords =>
      comparePerComponent(
        actTexelView.ulpFromZero(coords),
        expTexelView.ulpFromZero(coords),
        maxDiffULPsForNormFormat,
        altExpTexelView?.ulpFromZero(coords)
      );
  } else if (fmtIsFloat && maxDiffULPsForFloatFormat !== undefined) {
    tvc.predicate = coords =>
      comparePerComponent(
        actTexelView.ulpFromZero(coords),
        expTexelView.ulpFromZero(coords),
        maxDiffULPsForFloatFormat,
        altExpTexelView?.ulpFromZero(coords)
      );
  } else if (maxFractionalDiff !== undefined) {
    tvc.predicate = coords =>
      comparePerComponent(
        actTexelView.color(coords),
        expTexelView.color(coords),
        maxFractionalDiff,
        altExpTexelView?.color(coords)
      );
  } else {
    if (fmtIsNorm) {
      unreachable('need maxFractionalDiff or maxDiffULPsForNormFormat to compare norm textures');
    } else if (fmtIsFloat) {
      unreachable('need maxFractionalDiff or maxDiffULPsForFloatFormat to compare float textures');
    } else {
      unreachable();
    }
  }

  const repr = kTexelRepresentationInfo[format];
  if (fmtIsInt) {
    tvc.tableRows = failedCoords => [
      [`tolerance ± ${maxIntDiff}`],
      (function* () {
        yield* [` diff (act - exp)`, '==', ''];
        if (altExpTexelView) {
          yield* [` diff (act - alternateExp)`, '==', ''];
        }
        for (const coords of failedCoords) {
          const act = actTexelView.color(coords);
          const exp = expTexelView.color(coords);
          yield repr.componentOrder.map(ch => act[ch]! - exp[ch]!).join(',');
          if (altExpTexelView) {
            const altExp = altExpTexelView.color(coords);
            yield repr.componentOrder.map(ch => act[ch]! - altExp[ch]!).join(',');
          }
        }
      })(),
    ];
  } else if (
    (fmtIsNorm && maxDiffULPsForNormFormat !== undefined) ||
    (fmtIsFloat && maxDiffULPsForFloatFormat !== undefined)
  ) {
    const toleranceULPs = fmtIsNorm ? maxDiffULPsForNormFormat! : maxDiffULPsForFloatFormat!;
    tvc.tableRows = failedCoords => [
      [`tolerance ± ${toleranceULPs} normal-ULPs`],
      (function* () {
        yield* [` diff (act - exp) in normal-ULPs`, '==', ''];
        if (altExpTexelView) {
          yield* [` diff (act - alternateExp) in normal-ULPs`, '==', ''];
        }
        for (const coords of failedCoords) {
          const act = actTexelView.ulpFromZero(coords);
          const exp = expTexelView.ulpFromZero(coords);
          yield repr.componentOrder.map(ch => act[ch]! - exp[ch]!).join(',');
          if (altExpTexelView) {
            const altExp = altExpTexelView.ulpFromZero(coords);
            yield repr.componentOrder.map(ch => act[ch]! - altExp[ch]!).join(',');
          }
        }
      })(),
    ];
  } else {
    assert(maxFractionalDiff !== undefined);
    tvc.tableRows = failedCoords => [
      [`tolerance ± ${maxFractionalDiff}`],
      (function* () {
        yield* [` diff (act - exp)`, '==', ''];
        if (altExpTexelView) {
          yield* [` diff (act - alternateExp)`, '==', ''];
        }
        for (const coords of failedCoords) {
          const act = actTexelView.color(coords);
          const exp = expTexelView.color(coords);
          yield repr.componentOrder.map(ch => (act[ch]! - exp[ch]!).toPrecision(4)).join(',');
          if (altExpTexelView) {
            const altExp = altExpTexelView.color(coords);
            yield repr.componentOrder.map(ch => (act[ch]! - altExp[ch]!).toPrecision(4)).join(',');
          }
        }
      })(),
    ];
  }

  return tvc;
}

/** Returns true if all components of an "actual" texel compare as
 * nearly equal (with tolerance) against either a (required) expected
 * texel or against an optional alternate texel.
 */
function comparePerComponent(
  actual: PerTexelComponent<number>,
  expected: PerTexelComponent<number>,
  maxDiff: number,
  alternateExpected?: PerTexelComponent<number>
) {
  const similar = (a: PerTexelComponent<number>, b: PerTexelComponent<number>, maxDiff: number) =>
    Object.keys(a).every(key => {
      const k = key as TexelComponent;
      const act = a[k]!;
      const exp = b[k];
      if (exp === undefined) return false;
      return numbersApproximatelyEqual(act, exp, maxDiff);
    });

  if (similar(actual, expected, maxDiff)) {
    return true;
  }
  if (alternateExpected) {
    return similar(actual, alternateExpected, maxDiff);
  }
  return false;
}

/** Create a new mappable GPUBuffer, and copy a subrectangle of GPUTexture data into it. */
function createTextureCopyForMapRead(
  t: GPUTestBase,
  source: GPUTexelCopyTextureInfo,
  copySize: GPUExtent3D,
  { format }: { format: EncodableTextureFormat }
): { buffer: GPUBuffer; bytesPerRow: number; rowsPerImage: number } {
  const { byteLength, bytesPerRow, rowsPerImage } = getTextureSubCopyLayout(format, copySize, {
    aspect: source.aspect,
  });

  const buffer = t.createBufferTracked({
    usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
    size: byteLength,
  });

  const cmd = t.device.createCommandEncoder({ label: 'createTextureCopyForMapRead' });
  cmd.copyTextureToBuffer(source, { buffer, bytesPerRow, rowsPerImage }, copySize);
  t.device.queue.submit([cmd.finish()]);

  return { buffer, bytesPerRow, rowsPerImage };
}

export function findFailedPixels(
  format: EncodableTextureFormat,
  subrectOrigin: Required<GPUOrigin3DDict>,
  subrectSize: Required<GPUExtent3DDict>,
  {
    actTexelView,
    expTexelView,
    altExpTexelView,
  }: { actTexelView: TexelView; expTexelView: TexelView; altExpTexelView?: TexelView },
  texelCompareOptions: TexelCompareOptions,
  coords?: Generator<Required<GPUOrigin3DDict>>
) {
  const comparer = makeTexelViewComparer(
    format,
    { actTexelView, expTexelView, altExpTexelView },
    texelCompareOptions
  );

  const lowerCorner = [subrectSize.width, subrectSize.height, subrectSize.depthOrArrayLayers];
  const upperCorner = [0, 0, 0];
  const failedPixels: Required<GPUOrigin3DDict>[] = [];
  for (const coord of coords ?? fullSubrectCoordinates(subrectOrigin, subrectSize)) {
    const { x, y, z } = coord;
    if (!comparer.predicate(coord)) {
      failedPixels.push(coord);
      lowerCorner[0] = Math.min(lowerCorner[0], x);
      lowerCorner[1] = Math.min(lowerCorner[1], y);
      lowerCorner[2] = Math.min(lowerCorner[2], z);
      upperCorner[0] = Math.max(upperCorner[0], x);
      upperCorner[1] = Math.max(upperCorner[1], y);
      upperCorner[2] = Math.max(upperCorner[2], z);
    }
  }
  if (failedPixels.length === 0) {
    return undefined;
  }

  const repr = kTexelRepresentationInfo[format];
  // MAINTENANCE_TODO: Print depth-stencil formats as float+int instead of float+float.
  const printAsInteger = isColorTextureFormat(format)
    ? // For color, pick the type based on the format type
      ['uint', 'sint'].includes(getTextureFormatType(format))
    : // Print depth as "float", depth-stencil as "float,float", stencil as "int".
      !isDepthTextureFormat(format);
  const numericToString = numericToStringBuilder(printAsInteger);

  const componentOrderStr = repr.componentOrder.join(',') + ':';

  const printCoords = (function* () {
    yield* [' coords', '==', 'X,Y,Z:'];
    for (const coords of failedPixels) yield `${coords.x},${coords.y},${coords.z}`;
  })();
  const printActualBytes = (function* () {
    yield* [' act. texel bytes (little-endian)', '==', '0x:'];
    for (const coords of failedPixels) {
      yield Array.from(actTexelView.bytes(coords), b => b.toString(16).padStart(2, '0')).join(' ');
    }
  })();
  const printActualColors = (function* () {
    yield* [' act. colors', '==', componentOrderStr];
    for (const coords of failedPixels) {
      const pixel = actTexelView.color(coords);
      yield `${repr.componentOrder.map(ch => numericToString(pixel[ch]!)).join(',')}`;
    }
  })();
  const printExpectedColors = (function* () {
    yield* [' exp. colors', '==', componentOrderStr];
    for (const coords of failedPixels) {
      const pixel = expTexelView.color(coords);
      yield `${repr.componentOrder.map(ch => numericToString(pixel[ch]!)).join(',')}`;
    }
  })();
  const printAlternateExpectedColors = (function* () {
    if (altExpTexelView) {
      yield* [' alternateExp. colors', '==', componentOrderStr];
      for (const coords of failedPixels) {
        const pixel = altExpTexelView.color(coords);
        yield `${repr.componentOrder.map(ch => numericToString(pixel[ch]!)).join(',')}`;
      }
    }
  })();
  const printActualULPs = (function* () {
    yield* [' act. normal-ULPs-from-zero', '==', componentOrderStr];
    for (const coords of failedPixels) {
      const pixel = actTexelView.ulpFromZero(coords);
      yield `${repr.componentOrder.map(ch => pixel[ch]).join(',')}`;
    }
  })();
  const printExpectedULPs = (function* () {
    yield* [` exp. normal-ULPs-from-zero`, '==', componentOrderStr];
    for (const coords of failedPixels) {
      const pixel = expTexelView.ulpFromZero(coords);
      yield `${repr.componentOrder.map(ch => pixel[ch]).join(',')}`;
    }
  })();
  const printAlternateExpectedULPs = (function* () {
    if (altExpTexelView) {
      yield* [` alternateExp. normal-ULPs-from-zero`, '==', componentOrderStr];
      for (const coords of failedPixels) {
        const pixel = altExpTexelView.ulpFromZero(coords);
        yield `${repr.componentOrder.map(ch => pixel[ch]).join(',')}`;
      }
    }
  })();

  const opts = {
    fillToWidth: 120,
    numericToString,
  };
  return `\
 between ${lowerCorner} and ${upperCorner} inclusive:
${generatePrettyTable(opts, [
  printCoords,
  printActualBytes,
  printActualColors,
  printExpectedColors,
  printAlternateExpectedColors,
  printActualULPs,
  printExpectedULPs,
  printAlternateExpectedULPs,
  ...comparer.tableRows(failedPixels),
])}`;
}

/**
 * Check the contents of a GPUTexture by reading it back (with copyTextureToBuffer+mapAsync), then
 * comparing the data with the data in `expTexelView`.
 *
 * The actual and expected texture data are both converted to the "NormalULPFromZero" format,
 * which is a signed number representing how far the number is from zero, in ULPs, skipping
 * subnormal numbers (where ULP is defined for float, normalized, and integer formats).
 */
export async function textureContentIsOKByT2B(
  t: GPUTestBase,
  source: GPUTexelCopyTextureInfo,
  copySize_: GPUExtent3D,
  { expTexelView }: { expTexelView: TexelView },
  texelCompareOptions: TexelCompareOptions,
  coords?: Generator<Required<GPUOrigin3DDict>>
): Promise<ErrorWithExtra | undefined> {
  const subrectOrigin = reifyOrigin3D(source.origin ?? [0, 0, 0]);
  const subrectSize = reifyExtent3D(copySize_);
  const format = expTexelView.format;

  const { buffer, bytesPerRow, rowsPerImage } = createTextureCopyForMapRead(
    t,
    source,
    subrectSize,
    { format }
  );

  await buffer.mapAsync(GPUMapMode.READ);
  const data = new Uint8Array(buffer.getMappedRange());

  const texelViewConfig = {
    bytesPerRow,
    rowsPerImage,
    subrectOrigin,
    subrectSize,
  } as const;

  const actTexelView = TexelView.fromTextureDataByReference(format, data, texelViewConfig);

  const failedPixelsMessage = findFailedPixels(
    format,
    subrectOrigin,
    subrectSize,
    { actTexelView, expTexelView },
    texelCompareOptions,
    coords
  );

  if (failedPixelsMessage === undefined) {
    return undefined;
  }

  const msg = 'Texture level had unexpected contents:\n' + failedPixelsMessage;
  return new ErrorWithExtra(msg, () => ({
    expTexelView,
    // Make a new TexelView with a copy of the data so we can unmap the buffer (debug mode only).
    actTexelView: TexelView.fromTextureDataByReference(format, data.slice(), texelViewConfig),
  }));
}