/**
 * Dithering engine — palette building, dithering algorithms, and color separations.
 *
 * Ported from Python hitherdither-based service to pure TypeScript with Sharp.
 * Supports 11 algorithms (8 error diffusion + 3 ordered).
 */

import sharp from "sharp";

// ── Types ──

export type DitherAlgorithm =
	| "floyd-steinberg"
	| "atkinson"
	| "stucki"
	| "burkes"
	| "sierra"
	| "sierra-two-row"
	| "sierra-lite"
	| "jarvis-judice-ninke"
	| "bayer"
	| "ordered"
	| "cluster-dot";

export type PaletteMode = "auto" | "grayscale" | "spot" | "custom";

export interface DitherOptions {
	algorithm?: DitherAlgorithm;
	paletteMode?: PaletteMode;
	numColors?: number;
	customColors?: string[];
	threshold?: number;
	order?: number;
}

export interface DitherResult {
	buffer: Buffer;
	algorithm: DitherAlgorithm;
	paletteMode: PaletteMode;
	numColors: number;
	colorsUsed: string[];
	cached: boolean;
}

export interface SeparationResult {
	composite: Buffer;
	separations: Map<string, Buffer>;
	colors: string[];
}

// ── Constants ──

const MAX_DITHER_DIM = 512;
const MAX_CACHE = 200;

// ── Cache ──

const ditherCache = new Map<string, { buffer: Buffer; meta: Omit<DitherResult, "buffer" | "cached"> }>();
const separationCache = new Map<string, SeparationResult>();

function cacheKey(...parts: (string | number)[]): string {
	// Simple hash via string joining (no crypto needed for cache keys)
	return parts.join("|");
}

function evict<V>(cache: Map<string, V>, maxSize = MAX_CACHE) {
	if (cache.size <= maxSize) return;
	const iter = cache.keys();
	while (cache.size > maxSize) {
		const next = iter.next();
		if (next.done) break;
		cache.delete(next.value);
	}
}

// ── Color helpers ──

function hexToRgb(hex: string): [number, number, number] {
	const h = hex.replace(/^#/, "");
	return [parseInt(h.slice(0, 2), 16), parseInt(h.slice(2, 4), 16), parseInt(h.slice(4, 6), 16)];
}

function rgbToHex(r: number, g: number, b: number): string {
	return `${r.toString(16).padStart(2, "0")}${g.toString(16).padStart(2, "0")}${b.toString(16).padStart(2, "0")}`.toUpperCase();
}

function colorDistance(r1: number, g1: number, b1: number, r2: number, g2: number, b2: number): number {
	return (r1 - r2) ** 2 + (g1 - g2) ** 2 + (b1 - b2) ** 2;
}

function nearestColor(r: number, g: number, b: number, palette: [number, number, number][]): number {
	let minDist = Infinity;
	let idx = 0;
	for (let i = 0; i < palette.length; i++) {
		const d = colorDistance(r, g, b, palette[i][0], palette[i][1], palette[i][2]);
		if (d < minDist) {
			minDist = d;
			idx = i;
		}
	}
	return idx;
}

// ── Palette building ──

/**
 * Build a color palette from an image using median-cut quantization,
 * grayscale ramp, or custom colors.
 */
export async function buildPalette(
	inputBuffer: Buffer,
	mode: PaletteMode,
	numColors = 8,
	customColors?: string[],
): Promise<[number, number, number][]> {
	if (mode === "custom" && customColors?.length) {
		return customColors.map(hexToRgb);
	}

	if (mode === "grayscale") {
		const step = Math.floor(255 / Math.max(numColors - 1, 1));
		return Array.from({ length: numColors }, (_, i) => {
			const v = Math.min(i * step, 255);
			return [v, v, v] as [number, number, number];
		});
	}

	if (mode === "spot" && customColors?.length) {
		return customColors.slice(0, numColors).map(hexToRgb);
	}

	// Auto mode: extract dominant colors via Sharp + quantization
	const { data, info } = await sharp(inputBuffer)
		.resize(64, 64, { fit: "cover" })
		.removeAlpha()
		.raw()
		.toBuffer({ resolveWithObject: true });

	// Simple median-cut quantization
	const pixels: [number, number, number][] = [];
	for (let i = 0; i < data.length; i += 3) {
		pixels.push([data[i], data[i + 1], data[i + 2]]);
	}

	return medianCut(pixels, numColors);
}

/** Simple median-cut quantization. */
function medianCut(pixels: [number, number, number][], numColors: number): [number, number, number][] {
	type Box = [number, number, number][];

	const boxes: Box[] = [pixels];

	while (boxes.length < numColors) {
		// Find the box with the largest range
		let maxRange = -1;
		let maxIdx = 0;
		let splitChannel = 0;

		for (let i = 0; i < boxes.length; i++) {
			const box = boxes[i];
			for (let ch = 0; ch < 3; ch++) {
				let min = 255, max = 0;
				for (const p of box) {
					if (p[ch] < min) min = p[ch];
					if (p[ch] > max) max = p[ch];
				}
				const range = max - min;
				if (range > maxRange) {
					maxRange = range;
					maxIdx = i;
					splitChannel = ch;
				}
			}
		}

		if (maxRange <= 0) break;

		const box = boxes.splice(maxIdx, 1)[0];
		box.sort((a, b) => a[splitChannel] - b[splitChannel]);
		const mid = Math.floor(box.length / 2);
		boxes.push(box.slice(0, mid), box.slice(mid));
	}

	// Average each box to get palette colors
	return boxes.map((box) => {
		if (box.length === 0) return [0, 0, 0] as [number, number, number];
		let rSum = 0, gSum = 0, bSum = 0;
		for (const [r, g, b] of box) {
			rSum += r; gSum += g; bSum += b;
		}
		const n = box.length;
		return [Math.round(rSum / n), Math.round(gSum / n), Math.round(bSum / n)] as [number, number, number];
	});
}

// ── Error diffusion kernels ──

type DiffusionKernel = { offsets: [number, number, number][] }; // [dx, dy, weight]

const KERNELS: Record<string, DiffusionKernel> = {
	"floyd-steinberg": {
		offsets: [[1, 0, 7 / 16], [-1, 1, 3 / 16], [0, 1, 5 / 16], [1, 1, 1 / 16]],
	},
	atkinson: {
		offsets: [
			[1, 0, 1 / 8], [2, 0, 1 / 8],
			[-1, 1, 1 / 8], [0, 1, 1 / 8], [1, 1, 1 / 8],
			[0, 2, 1 / 8],
		],
	},
	stucki: {
		offsets: [
			[1, 0, 8 / 42], [2, 0, 4 / 42],
			[-2, 1, 2 / 42], [-1, 1, 4 / 42], [0, 1, 8 / 42], [1, 1, 4 / 42], [2, 1, 2 / 42],
			[-2, 2, 1 / 42], [-1, 2, 2 / 42], [0, 2, 4 / 42], [1, 2, 2 / 42], [2, 2, 1 / 42],
		],
	},
	burkes: {
		offsets: [
			[1, 0, 8 / 32], [2, 0, 4 / 32],
			[-2, 1, 2 / 32], [-1, 1, 4 / 32], [0, 1, 8 / 32], [1, 1, 4 / 32], [2, 1, 2 / 32],
		],
	},
	sierra: {
		offsets: [
			[1, 0, 5 / 32], [2, 0, 3 / 32],
			[-2, 1, 2 / 32], [-1, 1, 4 / 32], [0, 1, 5 / 32], [1, 1, 4 / 32], [2, 1, 2 / 32],
			[-1, 2, 2 / 32], [0, 2, 3 / 32], [1, 2, 2 / 32],
		],
	},
	"sierra-two-row": {
		offsets: [
			[1, 0, 4 / 16], [2, 0, 3 / 16],
			[-2, 1, 1 / 16], [-1, 1, 2 / 16], [0, 1, 3 / 16], [1, 1, 2 / 16], [2, 1, 1 / 16],
		],
	},
	"sierra-lite": {
		offsets: [[1, 0, 2 / 4], [0, 1, 1 / 4], [-1, 1, 1 / 4]],
	},
	"jarvis-judice-ninke": {
		offsets: [
			[1, 0, 7 / 48], [2, 0, 5 / 48],
			[-2, 1, 3 / 48], [-1, 1, 5 / 48], [0, 1, 7 / 48], [1, 1, 5 / 48], [2, 1, 3 / 48],
			[-2, 2, 1 / 48], [-1, 2, 3 / 48], [0, 2, 5 / 48], [1, 2, 3 / 48], [2, 2, 1 / 48],
		],
	},
};

// ── Dithering algorithms ──

function errorDiffusionDither(
	data: Uint8Array,
	width: number,
	height: number,
	palette: [number, number, number][],
	kernel: DiffusionKernel,
): Uint8Array {
	// Work with float copy to handle error accumulation
	const pixels = new Float32Array(data.length);
	for (let i = 0; i < data.length; i++) pixels[i] = data[i];

	const output = new Uint8Array(data.length);

	for (let y = 0; y < height; y++) {
		for (let x = 0; x < width; x++) {
			const idx = (y * width + x) * 3;
			const r = Math.max(0, Math.min(255, Math.round(pixels[idx])));
			const g = Math.max(0, Math.min(255, Math.round(pixels[idx + 1])));
			const b = Math.max(0, Math.min(255, Math.round(pixels[idx + 2])));

			const nearest = nearestColor(r, g, b, palette);
			const [nr, ng, nb] = palette[nearest];

			output[idx] = nr;
			output[idx + 1] = ng;
			output[idx + 2] = nb;

			const errR = r - nr;
			const errG = g - ng;
			const errB = b - nb;

			for (const [dx, dy, weight] of kernel.offsets) {
				const nx = x + dx;
				const ny = y + dy;
				if (nx >= 0 && nx < width && ny >= 0 && ny < height) {
					const nIdx = (ny * width + nx) * 3;
					pixels[nIdx] += errR * weight;
					pixels[nIdx + 1] += errG * weight;
					pixels[nIdx + 2] += errB * weight;
				}
			}
		}
	}

	return output;
}

function bayerDither(
	data: Uint8Array,
	width: number,
	height: number,
	palette: [number, number, number][],
	threshold: number,
	order: number,
): Uint8Array {
	const matrix = generateBayerMatrix(order);
	const matrixSize = matrix.length;
	const output = new Uint8Array(data.length);

	for (let y = 0; y < height; y++) {
		for (let x = 0; x < width; x++) {
			const idx = (y * width + x) * 3;
			const bayerVal = matrix[y % matrixSize][x % matrixSize];
			const factor = (bayerVal - 0.5) * threshold;

			const r = Math.max(0, Math.min(255, Math.round(data[idx] + factor)));
			const g = Math.max(0, Math.min(255, Math.round(data[idx + 1] + factor)));
			const b = Math.max(0, Math.min(255, Math.round(data[idx + 2] + factor)));

			const nearest = nearestColor(r, g, b, palette);
			output[idx] = palette[nearest][0];
			output[idx + 1] = palette[nearest][1];
			output[idx + 2] = palette[nearest][2];
		}
	}

	return output;
}

function generateBayerMatrix(order: number): number[][] {
	if (order <= 1) return [[0]];

	const size = 1 << order; // 2^order
	const matrix: number[][] = Array.from({ length: size }, () => new Array(size));

	// Generate recursively
	const base = [[0, 2], [3, 1]];

	function fill(mat: number[][], size: number): void {
		if (size === 2) {
			for (let y = 0; y < 2; y++)
				for (let x = 0; x < 2; x++)
					mat[y][x] = base[y][x];
			return;
		}

		const half = size / 2;
		const sub: number[][] = Array.from({ length: half }, () => new Array(half));
		fill(sub, half);

		for (let y = 0; y < half; y++) {
			for (let x = 0; x < half; x++) {
				const val = sub[y][x] * 4;
				mat[y][x] = val;
				mat[y][x + half] = val + 2;
				mat[y + half][x] = val + 3;
				mat[y + half][x + half] = val + 1;
			}
		}
	}

	fill(matrix, size);

	// Normalize to [0, 1]
	const total = size * size;
	for (let y = 0; y < size; y++)
		for (let x = 0; x < size; x++)
			matrix[y][x] = (matrix[y][x] + 0.5) / total;

	return matrix;
}

function clusterDotDither(
	data: Uint8Array,
	width: number,
	height: number,
	palette: [number, number, number][],
	threshold: number,
): Uint8Array {
	// 4x4 cluster-dot pattern
	const cluster = [
		[12, 5, 6, 13],
		[4, 0, 1, 7],
		[11, 3, 2, 8],
		[15, 10, 9, 14],
	];
	const output = new Uint8Array(data.length);

	for (let y = 0; y < height; y++) {
		for (let x = 0; x < width; x++) {
			const idx = (y * width + x) * 3;
			const cVal = (cluster[y % 4][x % 4] / 16 - 0.5) * threshold;

			const r = Math.max(0, Math.min(255, Math.round(data[idx] + cVal)));
			const g = Math.max(0, Math.min(255, Math.round(data[idx + 1] + cVal)));
			const b = Math.max(0, Math.min(255, Math.round(data[idx + 2] + cVal)));

			const nearest = nearestColor(r, g, b, palette);
			output[idx] = palette[nearest][0];
			output[idx + 1] = palette[nearest][1];
			output[idx + 2] = palette[nearest][2];
		}
	}

	return output;
}

// ── Main API ──

/**
 * Dither a design image with the specified algorithm and palette.
 * Results are cached by slug+params combination.
 */
export async function ditherDesign(
	inputBuffer: Buffer,
	slug: string,
	options: DitherOptions = {},
): Promise<DitherResult> {
	const {
		algorithm = "floyd-steinberg",
		paletteMode = "auto",
		numColors = 8,
		customColors,
		threshold = 64,
		order = 3,
	} = options;

	const key = cacheKey(slug, algorithm, paletteMode, numColors, (customColors || []).join(","), threshold, order);

	const cached = ditherCache.get(key);
	if (cached) {
		return { ...cached.meta, buffer: cached.buffer, cached: true };
	}

	// Downscale for performance
	const resized = await sharp(inputBuffer)
		.resize(MAX_DITHER_DIM, MAX_DITHER_DIM, { fit: "inside", withoutEnlargement: true })
		.removeAlpha()
		.raw()
		.toBuffer({ resolveWithObject: true });

	const { data, info } = resized;
	const { width, height } = info;

	// Build palette
	const palette = await buildPalette(inputBuffer, paletteMode, numColors, customColors);

	// Apply dithering
	let dithered: Uint8Array;

	if (algorithm === "bayer" || algorithm === "ordered") {
		dithered = bayerDither(new Uint8Array(data), width, height, palette, threshold, order);
	} else if (algorithm === "cluster-dot") {
		dithered = clusterDotDither(new Uint8Array(data), width, height, palette, threshold);
	} else {
		const kernel = KERNELS[algorithm] || KERNELS["floyd-steinberg"];
		dithered = errorDiffusionDither(new Uint8Array(data), width, height, palette, kernel);
	}

	// Convert back to PNG
	const buffer = await sharp(Buffer.from(dithered), {
		raw: { width, height, channels: 3 },
	})
		.png({ compressionLevel: 9 })
		.toBuffer();

	const colorsUsed = palette.map(([r, g, b]) => rgbToHex(r, g, b));

	const meta = { algorithm, paletteMode, numColors: colorsUsed.length, colorsUsed };
	ditherCache.set(key, { buffer, meta });
	evict(ditherCache);

	return { ...meta, buffer, cached: false };
}

/**
 * Generate color separations for screen printing.
 * Returns composite dithered image + individual per-color separation PNGs.
 */
export async function generateColorSeparations(
	inputBuffer: Buffer,
	slug: string,
	numColors = 4,
	algorithm: DitherAlgorithm = "floyd-steinberg",
	spotColors?: string[],
): Promise<SeparationResult> {
	const paletteMode: PaletteMode = spotColors?.length ? "spot" : "auto";
	const key = cacheKey("sep", slug, numColors, algorithm, (spotColors || []).join(","));

	const cached = separationCache.get(key);
	if (cached) return cached;

	// First dither the image
	const ditherResult = await ditherDesign(inputBuffer, `${slug}-sep`, {
		algorithm,
		paletteMode,
		numColors,
		customColors: spotColors,
	});

	const composite = ditherResult.buffer;

	// Get raw pixel data from the dithered result
	const { data, info } = await sharp(composite)
		.removeAlpha()
		.raw()
		.toBuffer({ resolveWithObject: true });

	const { width, height } = info;
	const palette = ditherResult.colorsUsed.map(hexToRgb);
	const separations = new Map<string, Buffer>();

	// Generate per-color separation images
	for (let pi = 0; pi < palette.length; pi++) {
		const [pr, pg, pb] = palette[pi];
		const hexColor = ditherResult.colorsUsed[pi];
		const sepData = new Uint8Array(width * height * 3);

		// White background
		sepData.fill(255);

		for (let y = 0; y < height; y++) {
			for (let x = 0; x < width; x++) {
				const idx = (y * width + x) * 3;
				const r = data[idx], g = data[idx + 1], b = data[idx + 2];

				// Match with tolerance
				if (
					Math.abs(r - pr) < 16 &&
					Math.abs(g - pg) < 16 &&
					Math.abs(b - pb) < 16
				) {
					sepData[idx] = pr;
					sepData[idx + 1] = pg;
					sepData[idx + 2] = pb;
				}
			}
		}

		const sepBuffer = await sharp(Buffer.from(sepData), {
			raw: { width, height, channels: 3 },
		})
			.png({ compressionLevel: 9 })
			.toBuffer();

		separations.set(hexColor, sepBuffer);
	}

	const result: SeparationResult = {
		composite,
		separations,
		colors: ditherResult.colorsUsed,
	};

	separationCache.set(key, result);
	evict(separationCache);

	return result;
}

/** List all supported dithering algorithms. */
export const ALGORITHMS: DitherAlgorithm[] = [
	"floyd-steinberg",
	"atkinson",
	"stucki",
	"burkes",
	"sierra",
	"sierra-two-row",
	"sierra-lite",
	"jarvis-judice-ninke",
	"bayer",
	"ordered",
	"cluster-dot",
];