import { pipeline, type Pipeline, env } from '@huggingface/transformers';

env.allowLocalModels = true; // allow cached weights without forcing downloads each boot

export type DepthMap = { data: Float32Array; width: number; height: number };

let depthPipelinePromise: Promise<Pipeline> | null = null;
let depthPipelineInitStartedAt = 0;

function readPositiveMsEnv(name: string, fallbackMs: number): number {
  const value = Number(process.env[name] ?? '');
  if (!Number.isFinite(value) || value <= 0) return fallbackMs;
  return Math.max(1_000, Math.round(value));
}

function timeoutError(label: string, timeoutMs: number): Error {
  return new Error(`${label} timed out after ${timeoutMs}ms`);
}

async function withTimeout<T>(promise: Promise<T>, timeoutMs: number, label: string): Promise<T> {
  if (!(timeoutMs > 0)) return promise;
  return await new Promise<T>((resolve, reject) => {
    const timer = setTimeout(() => reject(timeoutError(label, timeoutMs)), timeoutMs);
    promise.then(
      (value) => {
        clearTimeout(timer);
        resolve(value);
      },
      (err) => {
        clearTimeout(timer);
        reject(err);
      },
    );
  });
}

async function getDepthPipeline() {
  const initTimeoutMs = readPositiveMsEnv('DEPTH_PIPELINE_INIT_TIMEOUT_MS', 60_000);
  const staleAfterMs = readPositiveMsEnv('DEPTH_PIPELINE_STALE_RESET_MS', initTimeoutMs * 2);
  if (depthPipelinePromise && depthPipelineInitStartedAt > 0) {
    const ageMs = Date.now() - depthPipelineInitStartedAt;
    if (ageMs > staleAfterMs) {
      depthPipelinePromise = null;
      depthPipelineInitStartedAt = 0;
    }
  }
  if (!depthPipelinePromise) {
    depthPipelineInitStartedAt = Date.now();
    depthPipelinePromise = withTimeout(
      pipeline('depth-estimation', 'onnx-community/depth-anything-v2-small', {
        device: 'cpu',
      }),
      initTimeoutMs,
      'depth pipeline init',
    ).catch((err) => {
      depthPipelinePromise = null;
      depthPipelineInitStartedAt = 0;
      throw err;
    });
  }
  const pipe = await depthPipelinePromise;
  depthPipelineInitStartedAt = 0;
  return pipe;
}

function toFloatArray(input: unknown): Float32Array {
  if (input instanceof Float32Array) return input;
  if (ArrayBuffer.isView(input)) {
    const view = input as ArrayBufferView;
    const byteOffset = view.byteOffset ?? 0;
    const byteLength = view.byteLength ?? 0;
    const buffer = view.buffer;
    if (buffer && byteLength > 0 && byteLength % 4 === 0) {
      return new Float32Array(buffer.slice(byteOffset, byteOffset + byteLength));
    }
    const arrayLike = input as any;
    const length = Number(arrayLike?.length ?? 0);
    if (Number.isFinite(length) && length > 0) {
      const out = new Float32Array(length);
      for (let i = 0; i < length; i++) out[i] = Number(arrayLike[i] ?? 0);
      return out;
    }
    throw new Error('Unsupported depth tensor view');
  }
  if (Array.isArray(input)) return Float32Array.from(input as number[]);
  throw new Error('Unsupported depth tensor format');
}

function inferDimsFromShape(shape: unknown): { width: number; height: number } | null {
  const dims = Array.isArray(shape) ? shape.map((v) => Number(v)) : null;
  if (!dims || dims.length < 2 || !dims.every((v) => Number.isFinite(v) && v > 0)) return null;

  // Common formats:
  // - [H, W]
  // - [1, H, W]
  // - [1, 1, H, W]
  if (dims.length === 2) return { height: dims[0], width: dims[1] };
  if (dims.length === 3) return { height: dims[1], width: dims[2] };
  return { height: dims[dims.length - 2], width: dims[dims.length - 1] };
}

function sniffMimeType(buf: Buffer): string {
  if (buf.length >= 2 && buf[0] === 0xff && buf[1] === 0xd8) return 'image/jpeg';
  if (buf.length >= 8 && buf[0] === 0x89 && buf[1] === 0x50 && buf[2] === 0x4e && buf[3] === 0x47) return 'image/png';
  if (buf.length >= 6) {
    const header = buf.subarray(0, 6).toString('ascii');
    if (header === 'GIF87a' || header === 'GIF89a') return 'image/gif';
  }
  return 'application/octet-stream';
}

function extractSingleChannelPlane(
  raw: Float32Array,
  width: number,
  height: number,
): Float32Array {
  const area = width * height;
  if (raw.length === area) return raw;
  // If we got RGBA/RGB-like data, take the first channel per pixel.
  if (raw.length >= area * 3) {
    const out = new Float32Array(area);
    const stride = raw.length >= area * 4 ? 4 : 3;
    for (let i = 0, j = 0; j < area; j++, i += stride) out[j] = raw[i];
    return out;
  }
  // If the tensor still contains batch/channel dims, slice the first plane.
  if (raw.length >= area) return raw.subarray(0, area);
  throw new Error('Depth output data length is smaller than expected plane size');
}

export async function estimateDepth(imageBuffer: Buffer): Promise<DepthMap> {
  const inferTimeoutMs = readPositiveMsEnv('DEPTH_INFER_TIMEOUT_MS', 40_000);
  const depthEstimator = await getDepthPipeline();
  // Transformers.js expects a supported image input (RawImage|string|URL|Blob|Canvas).
  // In Node, wrap the bytes in a Blob to satisfy RawImage.read().
  const arrayBuffer = imageBuffer.buffer.slice(
    imageBuffer.byteOffset,
    imageBuffer.byteOffset + imageBuffer.byteLength,
  ) as ArrayBuffer;
  const blob = new Blob([arrayBuffer], { type: sniffMimeType(imageBuffer) });
  const result: any = await withTimeout(
    depthEstimator(blob),
    inferTimeoutMs,
    'depth inference',
  );

  // Depth-estimation pipeline usually returns:
  // - `predicted_depth`: Tensor (Float32)
  // - `depth`: RawImage (uint8 visualization)
  // Prefer the tensor for stable numeric processing.
  const depthOut = result?.predicted_depth ?? result?.depth ?? result;

  const explicitWidth = Number(depthOut?.width ?? 0);
  const explicitHeight = Number(depthOut?.height ?? 0);
  const shapeDims = inferDimsFromShape(depthOut?.shape ?? depthOut?.dims ?? depthOut?.tensor?.dims ?? null);
  const width = (explicitWidth > 0 ? explicitWidth : (shapeDims?.width ?? 0));
  const height = (explicitHeight > 0 ? explicitHeight : (shapeDims?.height ?? 0));
  if (!(width > 0 && height > 0)) {
    throw new Error('Depth output is missing width/height');
  }
  const rawData = toFloatArray(depthOut?.data ?? depthOut?.tensor?.data ?? depthOut);
  if (!rawData?.length) {
    throw new Error('Depth output has no data');
  }
  const data = extractSingleChannelPlane(rawData, width, height);
  return { data, width, height };
}