import { correctedDimensions, squareToQuad } from "./geometry";
import { adjustPixels, estimateDeskew } from "./pixels";
import {
  MAX_SOURCE_PIXELS,
  type PageAdjustments,
  type ScanPage,
} from "./types";

const MAX_WORKING_SOURCE_PIXELS = 24_000_000;

function canvas(width: number, height: number): HTMLCanvasElement {
  const element = document.createElement("canvas");
  element.width = width;
  element.height = height;
  return element;
}

function context(element: HTMLCanvasElement) {
  const value = element.getContext("2d", { willReadFrequently: true });
  if (!value) throw new Error("This browser could not create a 2D canvas.");
  return value;
}

export async function inspectImage(
  file: File,
): Promise<{ width: number; height: number }> {
  const bitmap = await createImageBitmap(file);
  try {
    if (bitmap.width * bitmap.height > MAX_SOURCE_PIXELS)
      throw new RangeError(
        `Decoded image exceeds ${MAX_SOURCE_PIXELS.toLocaleString()} pixels.`,
      );
    return { width: bitmap.width, height: bitmap.height };
  } finally {
    bitmap.close();
  }
}

function sampleBilinear(
  data: Uint8ClampedArray,
  width: number,
  height: number,
  x: number,
  y: number,
  channel: number,
): number {
  const clampedX = Math.max(0, Math.min(width - 1, x));
  const clampedY = Math.max(0, Math.min(height - 1, y));
  const x0 = Math.floor(clampedX);
  const y0 = Math.floor(clampedY);
  const x1 = Math.min(width - 1, x0 + 1);
  const y1 = Math.min(height - 1, y0 + 1);
  const fx = clampedX - x0;
  const fy = clampedY - y0;
  const top =
    (data[(y0 * width + x0) * 4 + channel] ?? 0) * (1 - fx) +
    (data[(y0 * width + x1) * 4 + channel] ?? 0) * fx;
  const bottom =
    (data[(y1 * width + x0) * 4 + channel] ?? 0) * (1 - fx) +
    (data[(y1 * width + x1) * 4 + channel] ?? 0) * fx;
  return top * (1 - fy) + bottom * fy;
}

function perspectiveCorrect(
  source: ImageData,
  outputWidth: number,
  outputHeight: number,
  adjustments: PageAdjustments,
): ImageData {
  const mapping = squareToQuad(adjustments.quad);
  const output = new ImageData(outputWidth, outputHeight);
  for (let y = 0; y < outputHeight; y += 1) {
    const normalizedY = outputHeight === 1 ? 0 : y / (outputHeight - 1);
    for (let x = 0; x < outputWidth; x += 1) {
      const normalizedX = outputWidth === 1 ? 0 : x / (outputWidth - 1);
      const sourcePoint = mapping(normalizedX, normalizedY);
      const sourceX = sourcePoint.x * (source.width - 1);
      const sourceY = sourcePoint.y * (source.height - 1);
      const outputIndex = (y * outputWidth + x) * 4;
      output.data[outputIndex] = sampleBilinear(
        source.data,
        source.width,
        source.height,
        sourceX,
        sourceY,
        0,
      );
      output.data[outputIndex + 1] = sampleBilinear(
        source.data,
        source.width,
        source.height,
        sourceX,
        sourceY,
        1,
      );
      output.data[outputIndex + 2] = sampleBilinear(
        source.data,
        source.width,
        source.height,
        sourceX,
        sourceY,
        2,
      );
      output.data[outputIndex + 3] = 255;
    }
  }
  output.data.set(adjustPixels(output.data, adjustments));
  return output;
}

function rotate(source: HTMLCanvasElement, degrees: number): HTMLCanvasElement {
  const normalized = ((degrees % 360) + 360) % 360;
  if (Math.abs(normalized) < 0.001) return source;
  const radians = (normalized * Math.PI) / 180;
  const width = Math.ceil(
    Math.abs(source.width * Math.cos(radians)) +
      Math.abs(source.height * Math.sin(radians)),
  );
  const height = Math.ceil(
    Math.abs(source.width * Math.sin(radians)) +
      Math.abs(source.height * Math.cos(radians)),
  );
  const output = canvas(width, height);
  const outputContext = context(output);
  outputContext.fillStyle = "white";
  outputContext.fillRect(0, 0, width, height);
  outputContext.translate(width / 2, height / 2);
  outputContext.rotate(radians);
  outputContext.drawImage(source, -source.width / 2, -source.height / 2);
  return output;
}

function boundPixels(
  source: HTMLCanvasElement,
  maximumPixels: number,
): HTMLCanvasElement {
  if (source.width * source.height <= maximumPixels) return source;
  const scale = Math.sqrt(maximumPixels / (source.width * source.height));
  const output = canvas(
    Math.max(16, Math.floor(source.width * scale)),
    Math.max(16, Math.floor(source.height * scale)),
  );
  const outputContext = context(output);
  outputContext.imageSmoothingEnabled = true;
  outputContext.imageSmoothingQuality = "high";
  outputContext.drawImage(source, 0, 0, output.width, output.height);
  return output;
}

export interface RenderedPage {
  canvas: HTMLCanvasElement;
  sourceScale: number;
  outputScale: number;
}

export async function renderPage(
  page: ScanPage,
  maximumPixels: number,
): Promise<RenderedPage> {
  const bitmap = await createImageBitmap(page.source);
  try {
    const sourceScale = Math.min(
      1,
      Math.sqrt(MAX_WORKING_SOURCE_PIXELS / (bitmap.width * bitmap.height)),
    );
    const sourceCanvas = canvas(
      Math.max(1, Math.round(bitmap.width * sourceScale)),
      Math.max(1, Math.round(bitmap.height * sourceScale)),
    );
    const sourceContext = context(sourceCanvas);
    sourceContext.drawImage(
      bitmap,
      0,
      0,
      sourceCanvas.width,
      sourceCanvas.height,
    );
    const dimensions = correctedDimensions(
      page.adjustments.quad,
      sourceCanvas.width,
      sourceCanvas.height,
      maximumPixels,
    );
    const corrected = canvas(dimensions.width, dimensions.height);
    context(corrected).putImageData(
      perspectiveCorrect(
        sourceContext.getImageData(
          0,
          0,
          sourceCanvas.width,
          sourceCanvas.height,
        ),
        dimensions.width,
        dimensions.height,
        page.adjustments,
      ),
      0,
      0,
    );
    const rotated = rotate(
      corrected,
      page.adjustments.quarterTurns * 90 + page.adjustments.deskew,
    );
    return {
      canvas: boundPixels(rotated, maximumPixels),
      sourceScale,
      outputScale: dimensions.scale,
    };
  } finally {
    bitmap.close();
  }
}

export async function detectDeskew(page: ScanPage): Promise<number> {
  const bitmap = await createImageBitmap(page.source);
  try {
    const scale = Math.min(1, 420 / Math.max(bitmap.width, bitmap.height));
    const work = canvas(
      Math.max(16, Math.round(bitmap.width * scale)),
      Math.max(16, Math.round(bitmap.height * scale)),
    );
    const workContext = context(work);
    workContext.drawImage(bitmap, 0, 0, work.width, work.height);
    return estimateDeskew(
      workContext.getImageData(0, 0, work.width, work.height).data,
      work.width,
      work.height,
    );
  } finally {
    bitmap.close();
  }
}

export function canvasBlob(
  source: HTMLCanvasElement,
  type: "image/png" | "image/jpeg",
  quality = 0.92,
): Promise<Blob> {
  return new Promise((resolve, reject) => {
    source.toBlob(
      (blob) =>
        blob
          ? resolve(blob)
          : reject(new Error("The browser could not encode the page.")),
      type,
      quality,
    );
  });
}