package com.foxx.androidcast.media;

/*********************************************************************
 * RgbToYuv420.java
 * Created at: Sat 16 May 2026 18:50:36 +0200
 * Updated at: Wed 20 May 2026 15:17:13 +0200 by Anton Afanasyeu <a.afanasieff@gmail.com>
 * Commit: 5d8e82d2e60a21fff3138d2a394ee4e8b4c6dcb8
 * class RgbToYuv420
 * Contributors:
 *    - Anton Afanasyeu <a.afanasieff@gmail.com> (4 commits, 159 lines)
 *    - Cursor Agent (project assistant)
 * Digest: SHA256	7050aefb284fec862ccdc6dbfc7f59720671fa25eeef2ccfc396e07ee2d9d68e
**********************************************************************/
import android.graphics.Bitmap;
import android.media.Image;

/**
 * ARGB to YUV420 planar (I420) or semi-planar (NV12) for MediaCodec buffer input.
 */
public final class RgbToYuv420 {
    public enum Layout {
        I420,
        NV12
    }

    private RgbToYuv420() {}

    public static byte[] fromBitmap(Bitmap bitmap, Layout layout) {
        int width = bitmap.getWidth();
        int height = bitmap.getHeight();
        int[] argb = new int[width * height];
        bitmap.getPixels(argb, 0, width, 0, 0, width, height);
        return fromArgb(argb, width, height, layout);
    }

    public static byte[] fromArgb(int[] argb, int width, int height, Layout layout) {
        if (layout == Layout.NV12) {
            return toNv12(argb, width, height);
        }
        return toI420(argb, width, height);
    }

    public static int bufferSize(int width, int height) {
        return width * height + (width * height) / 2;
    }

    /** Fills {@code scratch} with I420 from an RGBA {@link Image} (reuses scratch when large enough). */
    public static byte[] fromRgbaImage(Image image, byte[] scratch) {
        if (image == null) {
            return null;
        }
        int width = image.getWidth();
        int height = image.getHeight();
        int need = bufferSize(width, height);
        byte[] out = scratch != null && scratch.length >= need ? scratch : new byte[need];
        Image.Plane[] planes = image.getPlanes();
        if (planes == null || planes.length == 0) {
            return null;
        }
        Image.Plane plane = planes[0];
        java.nio.ByteBuffer buffer = plane.getBuffer();
        if (buffer == null) {
            return null;
        }
        int rowStride = plane.getRowStride();
        int pixelStride = plane.getPixelStride();
        int[] argb = new int[width * height];
        int offset = 0;
        for (int row = 0; row < height; row++) {
            int rowStart = row * rowStride;
            for (int col = 0; col < width; col++) {
                int pixel = buffer.getInt(rowStart + col * pixelStride);
                argb[offset++] = pixel;
            }
        }
        byte[] i420 = toI420(argb, width, height);
        if (out != i420) {
            System.arraycopy(i420, 0, out, 0, i420.length);
        }
        return out;
    }

    private static byte[] toI420(int[] argb, int width, int height) {
        int frameSize = width * height;
        int chromaW = (width + 1) / 2;
        int chromaH = (height + 1) / 2;
        int chromaLen = chromaW * chromaH;
        byte[] out = new byte[frameSize + chromaLen * 2];
        int uOffset = frameSize;
        int vOffset = frameSize + chromaLen;

        for (int j = 0; j < height; j++) {
            for (int i = 0; i < width; i++) {
                int color = argb[j * width + i];
                int r = (color >> 16) & 0xff;
                int g = (color >> 8) & 0xff;
                int b = color & 0xff;
                int y = yFromRgb(r, g, b);
                out[j * width + i] = (byte) y;
                if ((j & 1) == 0 && (i & 1) == 0) {
                    int u = uFromRgb(r, g, b);
                    int v = vFromRgb(r, g, b);
                    int cx = i >> 1;
                    int cy = j >> 1;
                    int cIndex = cy * chromaW + cx;
                    out[uOffset + cIndex] = (byte) u;
                    out[vOffset + cIndex] = (byte) v;
                }
            }
        }
        return out;
    }

    private static byte[] toNv12(int[] argb, int width, int height) {
        int frameSize = width * height;
        int chromaW = (width + 1) / 2;
        int chromaH = (height + 1) / 2;
        byte[] out = new byte[frameSize + chromaW * chromaH * 2];
        int uvOffset = frameSize;

        for (int j = 0; j < height; j++) {
            for (int i = 0; i < width; i++) {
                int color = argb[j * width + i];
                int r = (color >> 16) & 0xff;
                int g = (color >> 8) & 0xff;
                int b = color & 0xff;
                out[j * width + i] = (byte) yFromRgb(r, g, b);
                if ((j & 1) == 0 && (i & 1) == 0) {
                    int u = uFromRgb(r, g, b);
                    int v = vFromRgb(r, g, b);
                    int cx = i >> 1;
                    int cy = j >> 1;
                    int uvIndex = uvOffset + (cy * chromaW + cx) * 2;
                    out[uvIndex] = (byte) u;
                    out[uvIndex + 1] = (byte) v;
                }
            }
        }
        return out;
    }

    private static int yFromRgb(int r, int g, int b) {
        return clamp(((66 * r + 129 * g + 25 * b + 128) >> 8) + 16);
    }

    private static int uFromRgb(int r, int g, int b) {
        return clamp(((-38 * r - 74 * g + 112 * b + 128) >> 8) + 128);
    }

    private static int vFromRgb(int r, int g, int b) {
        return clamp(((112 * r - 94 * g - 18 * b + 128) >> 8) + 128);
    }

    private static int clamp(int v) {
        if (v < 0) {
            return 0;
        }
        if (v > 255) {
            return 255;
        }
        return v;
    }
}
