import {
  alignTo,
  assertBlockAlignedOrigin,
  assertBlockAlignedSize,
  blockInfo,
  bytesPerTexel,
  isCompressedFormat,
  physicalExtent,
  texelCopyLayout,
  toTextureData,
} from "./utils.js";
import { debugLog, trackResource, untrackResource } from "./debug.js";
import { nextId, shared } from "./shared.js";
import type { BufferData, GpuContext, GpuTexture } from "./types.js";

// Depth/stencil texel sizes are implementation-defined; typical values for
// memory pressure estimates. Exotic formats missing from the upload table
// estimate to 4.
const DEPTH_STENCIL_TEXEL_BYTES: Partial<Record<GPUTextureFormat, number>> = {
  stencil8: 1,
  depth16unorm: 2,
  depth24plus: 4,
  "depth24plus-stencil8": 5,
  depth32float: 4,
  "depth32float-stencil8": 5,
};

/** Size of one texel in bytes, estimating formats absent from the upload table. */
function estimatedBytesPerTexel(format: GPUTextureFormat): number {
  const depthStencilBytes = DEPTH_STENCIL_TEXEL_BYTES[format];
  if (depthStencilBytes !== undefined) return depthStencilBytes;
  try {
    return bytesPerTexel(format);
  } catch {
    return 4;
  }
}

/** Bytes one mip level occupies, compressed levels counted in whole blocks. */
function levelByteSize(
  format: GPUTextureFormat,
  width: number,
  height: number,
): number {
  const block = blockInfo(format);
  if (!block) return width * height * estimatedBytesPerTexel(format);
  // A level smaller than one block still occupies a whole one.
  return (
    Math.ceil(width / block.width) *
    Math.ceil(height / block.height) *
    block.bytes
  );
}

export interface TextureByteSizeOptions extends Partial<
  Pick<
    GPUTextureDescriptor,
    "format" | "mipLevelCount" | "sampleCount" | "size"
  >
> {
  width?: number;
  height?: number;
  /** Depth or array layer count. Falls back to `depthOrArrayLayers`. */
  depth?: number;
  depthOrArrayLayers?: number;
}

/**
 * Estimated memory footprint of a texture, in bytes — from an existing
 * {@link GpuTexture} or from a descriptor, so allocation can be budgeted before
 * anything is created.
 *
 * An estimate, not a guarantee: depth/stencil texel sizes and any driver-side
 * padding, alignment or metadata (HTILE, DCC, ...) are implementation-defined.
 */
export function textureByteSize(
  descriptor: GpuTexture | TextureByteSizeOptions,
): number {
  const { width = 1, height = 1 } = descriptor;
  const layers =
    descriptor.depthOrArrayLayers ??
    (descriptor as TextureByteSizeOptions).depth ??
    1;
  const mipLevelCount = descriptor.mipLevelCount ?? 1;
  const sampleCount = descriptor.sampleCount ?? 1;
  const format = descriptor.format ?? "rgba8unorm";

  let bytes = 0;
  for (let level = 0; level < mipLevelCount; level++) {
    bytes +=
      levelByteSize(
        format,
        Math.max(1, width >> level),
        Math.max(1, height >> level),
      ) *
      layers *
      sampleCount;
  }
  return bytes;
}

export type ExternalImageSource =
  | ImageBitmap
  | HTMLImageElement
  | HTMLVideoElement
  | VideoFrame
  | HTMLCanvasElement
  | OffscreenCanvas;

const sourceSize = (source: ExternalImageSource) =>
  source instanceof HTMLVideoElement
    ? { width: source.videoWidth, height: source.videoHeight }
    : source instanceof VideoFrame
      ? { width: source.codedWidth, height: source.codedHeight }
      : { width: source.width, height: source.height };

/**
 * BufferData is raw bytes/numbers; anything else in the data union is an image
 * source.
 */
const isImageData = (
  data: BufferData | ExternalImageSource | ExternalImageSource[],
): data is ExternalImageSource | ExternalImageSource[] =>
  Array.isArray(data)
    ? typeof data[0] === "object" && !Array.isArray(data[0])
    : !ArrayBuffer.isView(data) && !(data instanceof ArrayBuffer);

export interface CreateTextureOptions
  extends
    GPUObjectDescriptorBase,
    Partial<
      Pick<GPUTextureDescriptor, "dimension" | "sampleCount" | "mipLevelCount">
    >,
    Partial<GPUExtent3DDict>,
    Pick<GPUCopyExternalImageSourceInfo, "flipY">,
    Pick<GPUCopyExternalImageDestInfo, "colorSpace" | "premultipliedAlpha"> {
  /** Texture format. Required in raw WebGPU; defaults to `"rgba8unorm"` here. */
  format?: GPUTextureFormat;
  /**
   * Initial contents for mip level 0: raw data, or one or more
   * images/canvases/video frames — a single source makes a 2D texture, an array
   * of 6 with `viewDimension: "cube"` makes a cubemap, any other array makes a
   * 2D texture array (one source per layer).
   */
  data?: BufferData | ExternalImageSource | ExternalImageSource[];
  /**
   * Depth or array layer count. Falls back to `depthOrArrayLayers`, then to the
   * number of image sources in `data`.
   */
  depth?: number;
  /**
   * Allocate and generate the full mip chain from the uploaded data. For
   * explicit level counts, see `mipLevelCount`.
   */
  mipmap?: boolean;
  /**
   * Pre-computed mip chain, largest first, one entry per level — the way to mip
   * compressed formats, which cannot be rendered into and so cannot use
   * `mipmap`. Sets `mipLevelCount`, and each level's size is the base size
   * halved per level.
   */
  mipLevels?: BufferData[];
  /** Dimension of the cached default view, eg. "cube" for cubemaps. */
  viewDimension?: GPUTextureViewDimension;
  /**
   * Usage flags. Required in raw WebGPU; defaults to `TEXTURE_BINDING |
   * COPY_DST | RENDER_ATTACHMENT` here — set explicitly to replace (not merge
   * with) that default.
   */
  usage?: GPUTextureUsageFlags;
}

/** Number of mip levels in a full chain for the given dimensions. */
export const fullMipLevelCount = (width: number, height = 1): number =>
  1 + Math.floor(Math.log2(Math.max(width, height)));

/** Type guard for {@link GpuTexture} resource objects. */
export const isGpuTexture = (value: unknown): value is GpuTexture =>
  typeof value === "object" &&
  value !== null &&
  "id" in value &&
  "texture" in value;

const compressedMipmapsError = (format: GPUTextureFormat) =>
  new Error(
    `pex-gpu: cannot render mipmaps into compressed format "${format}"; supply a pre-computed chain as \`mipLevels\` instead`,
  );

/**
 * Create a texture. Initial contents can be raw data or image sources —
 * dimensions, layer count and array view dimension are inferred from images.
 *
 * ```js
 * const texture = gpu.createTexture(ctx, {
 *   width: 2,
 *   height: 1,
 *   data: [255, 255, 255, 255, 0, 0, 0, 255],
 * });
 * const depth = gpu.createTexture(ctx, {
 *   width,
 *   height,
 *   format: "depth24plus",
 * });
 * const photo = gpu.createTexture(ctx, {
 *   data: image,
 *   mipmap: true,
 *   flipY: true,
 * });
 * const cubemap = gpu.createTexture(ctx, {
 *   data: faces,
 *   viewDimension: "cube",
 * });
 * const compressed = gpu.createTexture(ctx, {
 *   width,
 *   height,
 *   format: "bc7-rgba-unorm",
 *   mipLevels: levels,
 * });
 * ```
 */
export function createTexture(
  ctx: GpuContext,
  options: CreateTextureOptions,
): GpuTexture {
  const images =
    options.data !== undefined && isImageData(options.data)
      ? Array.isArray(options.data)
        ? options.data
        : [options.data]
      : undefined;
  const imageSize = images ? sourceSize(images[0]!) : undefined;

  const format = options.format ?? "rgba8unorm";
  const width = options.width ?? imageSize?.width;
  if (width === undefined) {
    throw new Error(
      "pex-gpu: createTexture requires a width or image data to infer it from",
    );
  }
  const height = options.height ?? imageSize?.height ?? 1;
  const depthOrArrayLayers =
    options.depth ??
    options.depthOrArrayLayers ??
    (images && images.length > 1 ? images.length : 1);
  const mipLevelCount =
    options.mipLevelCount ??
    options.mipLevels?.length ??
    (options.mipmap ? fullMipLevelCount(width, height) : 1);
  const dimension = options.dimension ?? "2d";
  const viewDimension: GPUTextureViewDimension =
    options.viewDimension ??
    (images && Array.isArray(options.data) ? "2d-array" : undefined) ??
    (dimension === "2d" && depthOrArrayLayers > 1 ? "2d-array" : dimension);

  // Validated before allocating: throwing afterwards would leak the texture,
  // since the caller never receives a handle to dispose.
  assertBlockAlignedSize(format, width, height);
  if (options.mipmap && options.mipLevels) {
    throw new Error(
      "pex-gpu: createTexture takes either mipLevels or mipmap, not both",
    );
  }
  if (options.mipmap && isCompressedFormat(format)) {
    throw compressedMipmapsError(format);
  }

  const texture = ctx.device.createTexture({
    size: { width, height, depthOrArrayLayers },
    format,
    mipLevelCount,
    sampleCount: options.sampleCount ?? 1,
    dimension,
    // Compressed formats support only COPY_SRC/COPY_DST/TEXTURE_BINDING, so
    // the usual RENDER_ATTACHMENT default would fail validation for them.
    usage:
      options.usage ??
      (isCompressedFormat(format)
        ? GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.COPY_DST
        : GPUTextureUsage.TEXTURE_BINDING |
          GPUTextureUsage.COPY_DST |
          GPUTextureUsage.RENDER_ATTACHMENT),
    ...(options.label && { label: options.label }),
  });

  const id = nextId();
  let disposed = false;
  function dispose(): void {
    if (disposed) return;
    disposed = true;
    untrackResource(ctx.device, "textures", id);
    texture.destroy();
  }

  const result: GpuTexture = {
    texture,
    view: texture.createView({ dimension: viewDimension }),
    format,
    width,
    height,
    depthOrArrayLayers,
    mipLevelCount,
    sampleCount: options.sampleCount ?? 1,
    dimension,
    viewDimension,
    id,
    dispose,
    [Symbol.dispose]: dispose,
  };

  const bytes = textureByteSize(result);
  trackResource(ctx.device, "textures", {
    id,
    ...(options.label && { label: options.label }),
    usage: texture.usage,
    format,
    width,
    height,
    depthOrArrayLayers,
    mipLevelCount,
    sampleCount: result.sampleCount,
    dimension,
    bytes,
  });
  debugLog(
    ctx.device,
    "resources",
    () =>
      `createTexture ${options.label ? `"${options.label}" ` : ""}${width}×${height}` +
      `${depthOrArrayLayers > 1 ? `×${depthOrArrayLayers}` : ""} ${format}, ~${bytes} bytes`,
  );

  if (images) {
    for (const [layer, image] of images.entries()) {
      copyExternalImage(ctx, result, image, {
        origin: [0, 0, layer],
        ...(options.flipY !== undefined && { flipY: options.flipY }),
        ...(options.premultipliedAlpha !== undefined && {
          premultipliedAlpha: options.premultipliedAlpha,
        }),
        ...(options.colorSpace && { colorSpace: options.colorSpace }),
      });
    }
  } else if (options.mipLevels) {
    // Uploaded level by level: a compressed chain has no other route in, and a
    // supplied uncompressed one is authoritative over anything generated.
    for (const [mipLevel, data] of options.mipLevels.entries()) {
      updateTexture(ctx, result, data, { mipLevel });
    }
  } else if (options.data) {
    updateTexture(ctx, result, options.data as BufferData);
  }
  if (options.mipmap) generateMipmaps(ctx, result);
  return result;
}

export interface UpdateTextureOptions
  extends
    Pick<GPUTexelCopyTextureInfo, "origin" | "mipLevel">,
    Partial<Pick<GPUExtent3DDict, "width" | "height" | "depthOrArrayLayers">> {}

/**
 * Upload raw data into a texture region. Float formats accept a plain number
 * array (packed at the format's precision); integer, unorm, snorm and packed
 * formats require a matching typed array — its bytes are uploaded as-is.
 *
 * ```js
 * gpu.updateTexture(ctx, texture, pixels);
 * gpu.updateTexture(ctx, cubemap, facePixels, {
 *   origin: [0, 0, faceIndex],
 * });
 * ```
 */
export function updateTexture(
  ctx: GpuContext,
  target: GpuTexture,
  data: BufferData,
  options: UpdateTextureOptions = {},
): void {
  const mipLevel = options.mipLevel ?? 0;
  const width = options.width ?? Math.max(1, target.width >> mipLevel);
  const height = options.height ?? Math.max(1, target.height >> mipLevel);
  const typed = toTextureData(data, target.format);
  const origin = options.origin ?? [0, 0, 0];
  assertBlockAlignedOrigin(target.format, origin);

  // A compressed subresource is physically whole blocks even when its logical
  // size is smaller — the 2x2 tail of a 4x4-block mip chain still occupies one
  // 4x4 block — and copy extents must match that physical size, which
  // texelCopyLayout already measures the data in.
  const extent = physicalExtent(target.format, width, height);

  ctx.device.queue.writeTexture(
    { texture: target.texture, mipLevel, origin },
    typed as GPUAllowSharedBufferSource,
    texelCopyLayout(target.format, width, height),
    {
      ...extent,
      depthOrArrayLayers: options.depthOrArrayLayers ?? 1,
    },
  );
}

export interface CopyExternalImageOptions
  extends
    Pick<GPUCopyExternalImageSourceInfo, "flipY">,
    Pick<GPUTexelCopyTextureInfo, "origin" | "mipLevel">,
    Pick<GPUCopyExternalImageDestInfo, "colorSpace" | "premultipliedAlpha"> {
  /**
   * Copy extent override. Defaults to the source size, shifted down for mip
   * levels.
   */
  width?: number;
  height?: number;
}

/** Copy an image, canvas or video frame into a texture. */
export function copyExternalImage(
  ctx: GpuContext,
  target: GpuTexture,
  source: ExternalImageSource,
  options: CopyExternalImageOptions = {},
): void {
  const mipLevel = options.mipLevel ?? 0;
  const { width, height } = sourceSize(source);
  ctx.device.queue.copyExternalImageToTexture(
    { source, flipY: options.flipY ?? false },
    {
      texture: target.texture,
      mipLevel,
      origin: options.origin ?? [0, 0, 0],
      premultipliedAlpha: options.premultipliedAlpha ?? false,
      ...(options.colorSpace && { colorSpace: options.colorSpace }),
    },
    {
      width: options.width ?? Math.max(1, width >> mipLevel),
      height: options.height ?? Math.max(1, height >> mipLevel),
    },
  );
}

// ─── Mipmap generation ───────────────────────────────────────────────────────

const BLIT_WGSL = /* wgsl */ `
struct Out {
  @builtin(position) position: vec4<f32>,
  @location(0) uv: vec2<f32>,
}

@vertex
fn vertexMain(@builtin(vertex_index) index: u32) -> Out {
  // Fullscreen triangle from vertex_index alone — no vertex/index buffer needed.
  let uv = vec2<f32>(f32((index << 1u) & 2u), f32(index & 2u));
  return Out(vec4<f32>(uv * 2.0 - 1.0, 0.0, 1.0), vec2<f32>(uv.x, 1.0 - uv.y));
}

@group(0) @binding(0) var sourceSampler: sampler;
@group(0) @binding(1) var sourceTexture: texture_2d<f32>;

@fragment
fn fragmentMain(input: Out) -> @location(0) vec4<f32> {
  return textureSample(sourceTexture, sourceSampler, input.uv);
}
`;

interface MipmapPipelines {
  sampler: GPUSampler;
  pipelines: Map<GPUTextureFormat, GPURenderPipeline>;
}

const mipmapCache = shared(
  "mipmapCache",
  () => new WeakMap<GPUDevice, MipmapPipelines>(),
);

export interface GenerateMipmapsOptions {
  /**
   * Record into this encoder instead of creating one. Pass the encoder of an
   * open batched frame (e.g. from inside `gpu.frame` or a frame graph pass) to
   * keep the mip chain in submission order with the rest of that frame; the
   * caller owns `finish()`/`submit()` in that case. Omitted, a dedicated
   * encoder is created and submitted immediately.
   */
  encoder?: GPUCommandEncoder;
}

/**
 * Generate the full mip chain of a 2D or 2D-array texture by downsampling each
 * level with a linear-filtered blit.
 */
export function generateMipmaps(
  ctx: GpuContext,
  target: GpuTexture,
  options: GenerateMipmapsOptions = {},
): void {
  if (target.mipLevelCount < 2) return;
  if (target.dimension === "3d") {
    throw new Error("pex-gpu: generateMipmaps does not support 3D textures");
  }
  // Downsampling renders into each level, and compressed formats can never be
  // render attachments. Their mip chain has to come from the encoder instead.
  if (isCompressedFormat(target.format)) {
    throw compressedMipmapsError(target.format);
  }

  const cache = mipmapCache.getOrInsertComputed(ctx.device, () => ({
    sampler: ctx.device.createSampler({
      minFilter: "linear",
      magFilter: "linear",
    }),
    pipelines: new Map(),
  }));

  const pipeline = cache.pipelines.getOrInsertComputed(target.format, () => {
    const module = ctx.pipelineCache.getShaderModule(BLIT_WGSL);
    return ctx.device.createRenderPipeline({
      label: `pex-gpu mipmap ${target.format}`,
      layout: "auto",
      vertex: { module, entryPoint: "vertexMain" },
      fragment: {
        module,
        entryPoint: "fragmentMain",
        targets: [{ format: target.format }],
      },
      primitive: { topology: "triangle-list" },
    });
  });

  const encoder =
    options.encoder ??
    ctx.device.createCommandEncoder({ label: "pex-gpu mipmaps" });
  for (let layer = 0; layer < target.depthOrArrayLayers; layer++) {
    for (let level = 1; level < target.mipLevelCount; level++) {
      const bindGroup = ctx.device.createBindGroup({
        layout: pipeline.getBindGroupLayout(0),
        entries: [
          { binding: 0, resource: cache.sampler },
          {
            binding: 1,
            resource: target.texture.createView({
              dimension: "2d",
              baseMipLevel: level - 1,
              mipLevelCount: 1,
              baseArrayLayer: layer,
              arrayLayerCount: 1,
            }),
          },
        ],
      });
      const pass = encoder.beginRenderPass({
        colorAttachments: [
          {
            view: target.texture.createView({
              dimension: "2d",
              baseMipLevel: level,
              mipLevelCount: 1,
              baseArrayLayer: layer,
              arrayLayerCount: 1,
            }),
            loadOp: "clear",
            storeOp: "store",
          },
        ],
      });
      pass.setPipeline(pipeline);
      pass.setBindGroup(0, bindGroup);
      pass.draw(3);
      pass.end();
    }
  }
  if (!options.encoder) ctx.device.queue.submit([encoder.finish()]);
}

/**
 * Row pitch aligned to 256, the alignment `copyBufferToTexture` and
 * `copyTextureToBuffer` require (`queue.writeTexture` does not — see
 * {@link texelCopyLayout}). Compressed formats are measured in whole blocks.
 */
export const paddedBytesPerRow = (
  width: number,
  format: GPUTextureFormat,
): number => alignTo(texelCopyLayout(format, width, 1).bytesPerRow, 256);