import sharp from 'sharp';
import type { TileSegmentationOptions, TileSegmentationResult } from './tile-unet';
import { inferTileSegmentation } from './tile-unet';
import type { OccluderSegOptions } from './occluder';
import { inferOccluderBest } from './occluder_best';
import {
  combineTileMaskWithOccluder,
  type CombineTileMaskOptions,
  type OccluderMaskPayload,
  type TileMaskCombineResult,
} from './tile-mask';

export type FullSegmentationRequest = {
  buffer: Buffer;
  width: number;
  height: number;
  filename?: string | null;
  mimeType?: string | null;
};

export type FullSegmentationOccluderOptions = OccluderSegOptions & {
  workerUrl?: string | null;
  workerTimeoutMs?: number;
  preferWorker?: boolean;
};

export type FullSegmentationOptions = {
  tile?: TileSegmentationOptions;
  occluder?: FullSegmentationOccluderOptions;
  combine?: CombineTileMaskOptions;
};

export type FullSegmentationResult = {
  tile: TileSegmentationResult;
  occluder: OccluderMaskPayload | null;
  combined: TileMaskCombineResult['combined'];
  versions: {
    raw: string;
    combined: string;
    occluder?: string | null;
  };
};

function coverage(mask: Uint8Array | null): number {
  if (!mask || !mask.length) return 0;
  let solid = 0;
  for (let i = 0; i < mask.length; i++) {
    if (mask[i] > 127) solid += 1;
  }
  return solid / mask.length;
}

async function resizeMask(
  mask: Uint8Array,
  srcWidth: number,
  srcHeight: number,
  dstWidth: number,
  dstHeight: number,
): Promise<Uint8Array> {
  if (srcWidth === dstWidth && srcHeight === dstHeight) return mask;
  const resized = await sharp(Buffer.from(mask), {
    raw: { width: srcWidth, height: srcHeight, channels: 1 },
  })
    .resize(dstWidth, dstHeight, { kernel: 'nearest' })
    .toColorspace('b-w')
    .raw()
    .toBuffer();
  return new Uint8Array(resized);
}

async function encodeMask(mask: Uint8Array, width: number, height: number): Promise<string> {
  const buffer = await sharp(Buffer.from(mask), {
    raw: { width, height, channels: 1 },
  })
    .toColorspace('b-w')
    .png({ compressionLevel: 4 })
    .toBuffer();
  return `data:image/png;base64,${buffer.toString('base64')}`;
}

async function callOccluderWorker(
  req: FullSegmentationRequest,
  options: FullSegmentationOccluderOptions,
): Promise<OccluderMaskPayload | null> {
  if (!options.workerUrl) return null;
  const timeout = Math.max(1_000, options.workerTimeoutMs ?? 20_000);
  const form = new FormData();
  const blob = new Blob([req.buffer], { type: req.mimeType ?? 'application/octet-stream' });
  form.append('file', blob, req.filename ?? 'upload');
  const response = await fetch(options.workerUrl, {
    method: 'POST',
    body: form,
    signal: AbortSignal.timeout(timeout),
  });
  if (!response.ok) {
    throw new Error(`occluder worker responded ${response.status}`);
  }
  const payload = await response.json() as any;
  if (!payload?.ok || typeof payload.mask !== 'string') {
    throw new Error('occluder worker returned invalid payload');
  }
  const png = Buffer.from(payload.mask, 'base64');
  const { data, info } = await sharp(png, { failOnError: false })
    .ensureAlpha()
    .removeAlpha()
    .greyscale()
    .raw()
    .toBuffer({ resolveWithObject: true });

  const mask = new Uint8Array(data);
  const srcWidth = info.width ?? req.width;
  const srcHeight = info.height ?? req.height;
  const resized = await resizeMask(mask, srcWidth, srcHeight, req.width, req.height);

  return {
    data: resized,
    width: req.width,
    height: req.height,
    dataUrl: await encodeMask(resized, req.width, req.height),
    source: payload.source ?? 'occluder-worker',
    coverage: coverage(resized),
  };
}

async function inferOccluder(
  req: FullSegmentationRequest,
  options: FullSegmentationOccluderOptions | undefined,
): Promise<OccluderMaskPayload | null> {
  if (!options) return null;
  if (options.workerUrl && options.preferWorker !== false) {
    try {
      return await callOccluderWorker(req, options);
    } catch (err) {
      console.warn('[full-seg] occluder worker failed, trying local fallback', err);
    }
  }
  try {
    const result = await inferOccluderBest(req.buffer, {
      debugDir: options.debugDir,
      timeoutMs: options.timeoutMs,
      minCoverage: options.minCoverage,
    });
    const resized = await resizeMask(result.mask, result.width, result.height, req.width, req.height);
    return {
      data: resized,
      width: req.width,
      height: req.height,
      dataUrl: await encodeMask(resized, req.width, req.height),
      source: result.source,
      coverage: coverage(resized),
      meta: result.meta,
    };
  } catch (err) {
    console.warn('[full-seg] local occluder inference failed', err);
    return null;
  }
}

export async function runFullSegmentation(
  req: FullSegmentationRequest,
  options: FullSegmentationOptions = {},
): Promise<FullSegmentationResult> {
  const tile = await inferTileSegmentation(req.buffer, options.tile);
  const occluder = await inferOccluder(req, options.occluder);
  const combination = await combineTileMaskWithOccluder(tile, occluder, options.combine);

  return {
    tile,
    occluder,
    combined: combination.combined,
    versions: {
      raw: combination.rawDataUrl,
      combined: combination.combined.dataUrl,
      occluder: combination.occluderDataUrl,
    },
  };
}