import sharp from 'sharp';
import * as ort from 'onnxruntime-node';
import { readFile } from 'node:fs/promises';
import { dirname, resolve } from 'node:path';
import { fileURLToPath, pathToFileURL } from 'node:url';

const MODEL_DIR = resolve(dirname(fileURLToPath(import.meta.url)), '..');
const PADDING_MULTIPLE = 32;
export type ModelFormat = 'fp32' | 'fp16';
const MODEL_FILES: Record<ModelFormat, string> = {
    fp32: 'model.onnx',
    fp16: 'model.fp16.onnx',
};

export class Classifier {
    private constructor(
        private readonly session: ort.InferenceSession,
        private readonly preprocess: {
            resize_longest_side_px: number;
            mean: [number, number, number];
            std: [number, number, number];
        },
        private readonly labels: {
            screen: { labels: string[] };
            safety: { labels: string[] };
        },
    ) {}

    /**
     * Load the exported ONNX classifier and sidecar metadata.
     *
     * @param modelDir Directory containing the exported model artifacts.
     * @returns Ready-to-use classifier instance.
     */
    static async create(modelDir = MODEL_DIR, modelFormat: ModelFormat = 'fp32'): Promise<Classifier> {
        const session = await ort.InferenceSession.create(resolve(modelDir, 'onnx', MODEL_FILES[modelFormat]), {
            executionProviders: ['cpu'],
        });
        const preprocess = JSON.parse(await readFile(resolve(modelDir, 'preprocess.json'), 'utf8')) as {
            resize_longest_side_px: number;
            mean: [number, number, number];
            std: [number, number, number];
        };
        const labels = JSON.parse(await readFile(resolve(modelDir, 'inference/labels.json'), 'utf8')) as {
            screen: { labels: string[] };
            safety: { labels: string[] };
        };
        return new Classifier(session, preprocess, labels);
    }

    /**
     * Classify one image.
     *
     * @param imagePath Path to an image file.
     * @returns Predicted public screen and safety labels.
     */
    async classify(imagePath: string): Promise<{ screen: string; safety: string }> {
        const [prediction] = await this.classifyBatch([imagePath]);
        if (prediction === undefined) {
            throw new Error('No prediction was produced.');
        }
        return prediction;
    }

    /**
     * Classify a batch of images.
     *
     * @param imagePaths Paths to image files.
     * @returns Predicted public screen and safety labels for each image.
     */
    async classifyBatch(imagePaths: readonly string[]): Promise<Array<{ screen: string; safety: string }>> {
        if (imagePaths.length === 0) {
            return [];
        }
        const images = await Promise.all(imagePaths.map((imagePath) => preprocessImage(imagePath, this.preprocess)));
        const output = await this.session.run({ image: collateImages(images) });
        return decodePredictions(outputTensor(output, 'screen'), outputTensor(output, 'safety'), this.labels);
    }
}

/**
 * Classify one image with a classifier loaded from disk.
 *
 * @param imagePath Path to an image file.
 * @param modelDir Directory containing the exported model artifacts.
 * @returns Predicted public screen and safety labels.
 */
export async function classify(
    imagePath: string,
    modelDir = MODEL_DIR,
    modelFormat: ModelFormat = 'fp32',
): Promise<{ screen: string; safety: string }> {
    const classifier = await Classifier.create(modelDir, modelFormat);
    if (imagePath.length === 0) {
        throw new Error('Image path must not be empty.');
    }
    return classifier.classify(imagePath);
}

/**
 * Classify a batch of images with a classifier loaded from disk.
 *
 * @param imagePaths Paths to image files.
 * @param modelDir Directory containing the exported model artifacts.
 * @returns Predicted public screen and safety labels for each image.
 */
export async function classifyBatch(
    imagePaths: readonly string[],
    modelDir = MODEL_DIR,
    modelFormat: ModelFormat = 'fp32',
): Promise<Array<{ screen: string; safety: string }>> {
    const classifier = await Classifier.create(modelDir, modelFormat);
    if (imagePaths.length === 0) {
        return [];
    }
    return classifier.classifyBatch(imagePaths);
}

async function preprocessImage(
    imagePath: string,
    preprocess: {
        resize_longest_side_px: number;
        mean: [number, number, number];
        std: [number, number, number];
    },
): Promise<{ data: Float32Array; height: number; width: number }> {
    const { data, info: metadata } = await sharp(imagePath)
        .rotate()
        .resize({
            width: preprocess.resize_longest_side_px,
            height: preprocess.resize_longest_side_px,
            fit: 'inside',
            kernel: 'cubic',
        })
        .removeAlpha()
        .toColourspace('srgb')
        .raw()
        .toBuffer({ resolveWithObject: true });

    if (metadata.channels !== 3) {
        throw new Error(`Expected RGB image data, got ${String(metadata.channels)} channels.`);
    }

    const pixels = metadata.width * metadata.height;
    const tensor = new Float32Array(3 * pixels);

    for (let pixel = 0; pixel < pixels; pixel += 1) {
        const offset = pixel * 3;
        tensor[pixel] = (channel(data, offset) / 255 - preprocess.mean[0]) / preprocess.std[0];
        tensor[pixels + pixel] = (channel(data, offset + 1) / 255 - preprocess.mean[1]) / preprocess.std[1];
        tensor[2 * pixels + pixel] = (channel(data, offset + 2) / 255 - preprocess.mean[2]) / preprocess.std[2];
    }

    return {
        data: tensor,
        height: metadata.height,
        width: metadata.width,
    };
}

function collateImages(
    images: ReadonlyArray<{ data: Float32Array; height: number; width: number }>,
): ort.TypedTensor<'float32'> {
    const height = Math.ceil(Math.max(...images.map((image) => image.height)) / PADDING_MULTIPLE) * PADDING_MULTIPLE;
    const width = Math.ceil(Math.max(...images.map((image) => image.width)) / PADDING_MULTIPLE) * PADDING_MULTIPLE;
    const imagePixels = height * width;
    const tensor = new Float32Array(images.length * 3 * imagePixels);

    for (let imageIndex = 0; imageIndex < images.length; imageIndex += 1) {
        const image = imageAt(images, imageIndex);
        const sourcePixels = image.height * image.width;
        for (let channelIndex = 0; channelIndex < 3; channelIndex += 1) {
            for (let row = 0; row < image.height; row += 1) {
                const sourceStart = channelIndex * sourcePixels + row * image.width;
                const targetStart = imageIndex * 3 * imagePixels + channelIndex * imagePixels + row * width;
                tensor.set(image.data.subarray(sourceStart, sourceStart + image.width), targetStart);
            }
        }
    }

    return new ort.Tensor('float32', tensor, [images.length, 3, height, width]);
}

function decodePredictions(
    screen: ort.TypedTensor<'float32'>,
    safety: ort.TypedTensor<'float32'>,
    labels: {
        screen: { labels: string[] };
        safety: { labels: string[] };
    },
): Array<{ screen: string; safety: string }> {
    // ONNX emits flat screen and safety logits.
    const batchSize = tensorDim(screen, 0);
    const screenClassCount = tensorDim(screen, 1);
    const safetyBatchSize = tensorDim(safety, 0);
    const safetyClassCount = tensorDim(safety, 1);
    if (safetyBatchSize !== batchSize) {
        throw new Error(
            `Safety batch size ${String(safetyBatchSize)} does not match screen batch size ${String(batchSize)}.`,
        );
    }

    const predictions: Array<{ screen: string; safety: string }> = [];
    for (let row = 0; row < batchSize; row += 1) {
        predictions.push({
            screen: labelAt(labels.screen.labels, topIndex(screen.data, row * screenClassCount, screenClassCount)),
            safety: labelAt(labels.safety.labels, topIndex(safety.data, row * safetyClassCount, safetyClassCount)),
        });
    }
    return predictions;
}

function outputTensor(output: ort.InferenceSession.ReturnType, name: string): ort.TypedTensor<'float32'> {
    const tensor = output[name];
    if (tensor === undefined) {
        throw new Error(`Missing ONNX output "${name}".`);
    }
    if (tensor.type !== 'float32') {
        throw new Error(`Expected "${name}" to be float32, got ${tensor.type}.`);
    }
    return tensor as ort.TypedTensor<'float32'>;
}

function tensorDim(tensor: ort.TypedTensor<'float32'>, index: number): number {
    const dimension = tensor.dims[index];
    if (typeof dimension !== 'number') {
        throw new Error(`Missing tensor dimension ${String(index)}.`);
    }
    return dimension;
}

function labelAt(labels: readonly string[], index: number): string {
    const label = labels[index];
    if (label === undefined) {
        throw new Error(`Missing label for class ${String(index)}.`);
    }
    return label;
}

function topIndex(data: Float32Array, offset: number, count: number): number {
    let top = 0;
    let score = valueAt(data, offset);
    for (let index = 1; index < count; index += 1) {
        const candidate = valueAt(data, offset + index);
        if (candidate > score) {
            top = index;
            score = candidate;
        }
    }
    return top;
}

function imageAt(
    images: ReadonlyArray<{ data: Float32Array; height: number; width: number }>,
    index: number,
): { data: Float32Array; height: number; width: number } {
    const image = images[index];
    if (image === undefined) {
        throw new Error(`Missing image at index ${String(index)}.`);
    }
    return image;
}

function valueAt(data: Float32Array, index: number): number {
    const value = data[index];
    if (value === undefined) {
        throw new Error(`Missing tensor value at offset ${String(index)}.`);
    }
    return value;
}

function channel(data: Buffer, index: number): number {
    const channelByte = data[index];
    if (channelByte === undefined) {
        throw new Error(`Missing image channel at offset ${String(index)}.`);
    }
    return channelByte;
}

const script = process.argv[1];
if (script !== undefined && import.meta.url === pathToFileURL(resolve(script)).href) {
    const args = process.argv.slice(2);
    const modelFlag = args.indexOf('--model-dir');
    const modelDir = modelFlag >= 0 ? args[modelFlag + 1] : MODEL_DIR;
    if (modelFlag >= 0) {
        args.splice(modelFlag, 2);
    }
    const formatFlag = args.indexOf('--model-format');
    const modelFormat = formatFlag >= 0 ? parseModelFormat(args[formatFlag + 1]) : 'fp32';
    if (formatFlag >= 0) {
        args.splice(formatFlag, 2);
    }
    if (args.length === 0 || modelDir === undefined) {
        throw new Error('Usage: bun inference/typescript.ts <image...> [--model-dir export_dir] [--model-format fp32|fp16]');
    }
    const predictions = await classifyBatch(args, modelDir, modelFormat);
    process.stdout.write(`${JSON.stringify(predictions.length === 1 ? predictions[0] : predictions, null, 2)}\n`);
}

function parseModelFormat(value: string | undefined): ModelFormat {
    if (value === 'fp32' || value === 'fp16') {
        return value;
    }
    throw new Error('Model format must be fp32 or fp16.');
}