package dev.photosync.core.thumbnail;

import java.util.Arrays;
import java.util.Base64;
import java.util.Optional;

/**
 * Decoder for the ThumbHash placeholders Immich attaches to every asset.
 *
 * <p>A ThumbHash is around 25 bytes and already travels with the timeline
 * response, so the blurred stand-in for a tile costs no extra request -- the
 * grid can be fully painted the instant a month's metadata lands, and the real
 * thumbnails fade in behind it. That is what makes fast scrolling feel like the
 * web gallery this browser is imitating rather than a wall of grey boxes.
 *
 * <p>Only decoding lives here; the mod never produces hashes. The format is
 * Evan Wallace's, and the maths below deliberately mirrors the reference
 * implementation step for step so it can be diffed against it.
 */
public final class ThumbHash {

    /** The reference decoder's output is normalised to fit in a 32x32 box. */
    private static final int MAX_EDGE = 32;

    private ThumbHash() {
    }

    /** Decodes the base64 form carried in the timeline response. */
    public static Optional<ThumbImage> decode(String base64) {
        if (base64 == null || base64.isBlank()) {
            return Optional.empty();
        }
        try {
            return decode(Base64.getDecoder().decode(base64.trim()));
        } catch (IllegalArgumentException e) {
            return Optional.empty();
        }
    }

    /**
     * Decodes the raw hash.
     *
     * <p>Returns empty rather than throwing for anything malformed: a bad
     * placeholder should cost the player a grey tile, never a crashed screen.
     */
    public static Optional<ThumbImage> decode(byte[] hash) {
        if (hash == null || hash.length < 5) {
            return Optional.empty();
        }
        try {
            return Optional.of(decodeChecked(hash));
        } catch (ArrayIndexOutOfBoundsException e) {
            return Optional.empty();
        }
    }

    private static ThumbImage decodeChecked(byte[] hash) {
        int header24 = byteAt(hash, 0) | (byteAt(hash, 1) << 8) | (byteAt(hash, 2) << 16);
        int header16 = byteAt(hash, 3) | (byteAt(hash, 4) << 8);

        float lDc = (header24 & 63) / 63f;
        float pDc = ((header24 >> 6) & 63) / 31.5f - 1f;
        float qDc = ((header24 >> 12) & 63) / 31.5f - 1f;
        float lScale = ((header24 >> 18) & 31) / 31f;
        boolean hasAlpha = (header24 >>> 23) != 0;
        float pScale = ((header16 >> 3) & 63) / 63f;
        float qScale = ((header16 >> 9) & 63) / 63f;
        boolean landscape = (header16 >>> 15) != 0;

        int lx = Math.max(3, landscape ? (hasAlpha ? 5 : 7) : (header16 & 7));
        int ly = Math.max(3, landscape ? (header16 & 7) : (hasAlpha ? 5 : 7));

        float aDc = hasAlpha ? (byteAt(hash, 5) & 15) / 15f : 1f;
        float aScale = hasAlpha ? (byteAt(hash, 5) >> 4) / 15f : 0f;

        // Saturation is boosted by 1.25 on the chroma channels to undo the loss
        // from quantising them into four bits, exactly as the encoder expects.
        int[] cursor = {0};
        int acStart = hasAlpha ? 6 : 5;
        float[] lAc = readChannel(hash, acStart, cursor, lx, ly, lScale);
        float[] pAc = readChannel(hash, acStart, cursor, 3, 3, pScale * 1.25f);
        float[] qAc = readChannel(hash, acStart, cursor, 3, 3, qScale * 1.25f);
        float[] aAc = hasAlpha ? readChannel(hash, acStart, cursor, 5, 5, aScale) : new float[0];

        // The aspect ratio comes from the unclamped counts, matching the encoder.
        float ratio = (float) (landscape ? (hasAlpha ? 5 : 7) : (header16 & 7))
                / (landscape ? (header16 & 7) : (hasAlpha ? 5 : 7));
        int width = Math.max(1, Math.round(ratio > 1 ? MAX_EDGE : MAX_EDGE * ratio));
        int height = Math.max(1, Math.round(ratio > 1 ? MAX_EDGE / ratio : MAX_EDGE));

        // The reference recomputes these cosines per pixel; hoisting them out
        // turns the inner loop into pure multiply-add.
        int fxCount = Math.max(lx, hasAlpha ? 5 : 3);
        int fyCount = Math.max(ly, hasAlpha ? 5 : 3);
        float[][] fxTable = cosineTable(width, fxCount);
        float[][] fyTable = cosineTable(height, fyCount);

        int[] argb = new int[width * height];
        for (int y = 0; y < height; y++) {
            float[] fy = fyTable[y];
            for (int x = 0; x < width; x++) {
                float[] fx = fxTable[x];
                float l = lDc;
                float p = pDc;
                float q = qDc;
                float a = aDc;

                for (int cy = 0, j = 0; cy < ly; cy++) {
                    for (int cx = cy > 0 ? 0 : 1; cx * ly < lx * (ly - cy); cx++, j++) {
                        l += lAc[j] * fx[cx] * fy[cy] * 2f;
                    }
                }
                for (int cy = 0, j = 0; cy < 3; cy++) {
                    for (int cx = cy > 0 ? 0 : 1; cx < 3 - cy; cx++, j++) {
                        float f = fx[cx] * fy[cy] * 2f;
                        p += pAc[j] * f;
                        q += qAc[j] * f;
                    }
                }
                if (hasAlpha) {
                    for (int cy = 0, j = 0; cy < 5; cy++) {
                        for (int cx = cy > 0 ? 0 : 1; cx < 5 - cy; cx++, j++) {
                            a += aAc[j] * fx[cx] * fy[cy] * 2f;
                        }
                    }
                }

                float blue = l - 2f / 3f * p;
                float red = (3f * l - blue + q) / 2f;
                float green = red - q;
                argb[y * width + x] = (channel(a) << 24) | (channel(red) << 16)
                        | (channel(green) << 8) | channel(blue);
            }
        }
        return new ThumbImage(width, height, argb);
    }

    /**
     * Reads one channel's AC coefficients, four bits each, from a nibble stream
     * shared by all channels -- hence the caller-owned cursor.
     */
    private static float[] readChannel(byte[] hash, int start, int[] cursor, int nx, int ny, float scale) {
        float[] ac = new float[nx * ny];
        int count = 0;
        for (int cy = 0; cy < ny; cy++) {
            for (int cx = cy > 0 ? 0 : 1; cx * ny < nx * (ny - cy); cx++) {
                int index = cursor[0]++;
                int nibble = (byteAt(hash, start + (index >> 1)) >> ((index & 1) << 2)) & 15;
                ac[count++] = (nibble / 7.5f - 1f) * scale;
            }
        }
        return Arrays.copyOf(ac, count);
    }

    private static float[][] cosineTable(int size, int terms) {
        float[][] table = new float[size][terms];
        for (int i = 0; i < size; i++) {
            for (int c = 0; c < terms; c++) {
                table[i][c] = (float) Math.cos(Math.PI / size * (i + 0.5) * c);
            }
        }
        return table;
    }

    private static int channel(float value) {
        return Math.max(0, Math.min(255, Math.round(value * 255f)));
    }

    private static int byteAt(byte[] hash, int index) {
        return hash[index] & 0xFF;
    }
}
