package com.xypower.mppreview.utils;

import static com.xypower.mppreview.utils.HdrUtil.generateTimestamp;

import android.graphics.Bitmap;
import android.graphics.ImageFormat;
import android.media.Image;
import android.os.Environment;
import android.renderscript.Allocation;
import android.renderscript.Element;
import android.renderscript.RenderScript;
import android.renderscript.ScriptIntrinsicYuvToRGB;
import android.renderscript.Type;
import android.util.Log;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;

public class ImageConverterUtil {
    private  final String TAG = "ImageConverterUtil";

    public void saveYuvImageFromImageReader(Image image, RenderScript rs, String i) {
//        Image image = null;
        try {
//            // Acquire the latest image from ImageReader
//            image = imageReader.acquireLatestImage();
//            if (image == null) {
//                Log.e(TAG, "No image available");
//                return;
//            }

            // Convert Image to Bitmap directly using RenderScript
            Bitmap bitmap = imageYuvToBitmap(image, rs);

            // Save the bitmap to file (PNG format to avoid lossy compression)
            saveBitmapToFile(bitmap, i);

            Log.d(TAG, "Image saved successfully");

        } catch (Exception e) {
            Log.e(TAG, "Error saving image: " + e.getMessage());
        } finally {
            if (image != null) {
                image.close();
            }
        }
    }

    public Bitmap imageYuvToBitmap(Image image, RenderScript rs) {
        if (image.getFormat() != ImageFormat.YUV_420_888) {
            throw new IllegalArgumentException("Only YUV_420_888 format is supported");
        }

        int width = image.getWidth();
        int height = image.getHeight();

        // Convert YUV to NV21 format
        byte[] nv21Data = yuv420ToNv21(image);

        // Create output bitmap
        Bitmap bitmap = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888);

        // Use RenderScript for YUV to RGB conversion (more efficient)
        Allocation inputAllocation = Allocation.createSized(rs, Element.U8(rs), nv21Data.length);
        inputAllocation.copyFrom(nv21Data);

        Type.Builder yuvType = new Type.Builder(rs, Element.U8(rs))
                .setX(width)
                .setY(height)
                .setYuvFormat(ImageFormat.NV21);
        Allocation yuvAllocation = Allocation.createTyped(rs, yuvType.create(), Allocation.USAGE_SCRIPT);
        yuvAllocation.copyFrom(nv21Data);

        Allocation outputAllocation = Allocation.createFromBitmap(rs, bitmap);

        ScriptIntrinsicYuvToRGB yuvToRgbScript = ScriptIntrinsicYuvToRGB.create(rs, Element.U8_4(rs));
        yuvToRgbScript.setInput(yuvAllocation);
        yuvToRgbScript.forEach(outputAllocation);

        outputAllocation.copyTo(bitmap);

        // Clean up allocations
        inputAllocation.destroy();
        yuvAllocation.destroy();
        outputAllocation.destroy();
        yuvToRgbScript.destroy();

        return bitmap;
    }

    private byte[] yuv420ToNv21(Image image) {
        Image.Plane[] planes = image.getPlanes();
        ByteBuffer yBuffer = planes[0].getBuffer();
        ByteBuffer uBuffer = planes[1].getBuffer();
        ByteBuffer vBuffer = planes[2].getBuffer();

        int ySize = yBuffer.remaining();
        int width = image.getWidth();
        int height = image.getHeight();
        int uvSize = width * height / 4; // U and V are quarter size

        byte[] nv21 = new byte[ySize + uvSize * 2];

        // Copy Y plane as-is
        yBuffer.get(nv21, 0, ySize);

        // Interleave V and U planes into NV21 format (which is really NV12 with U and V swapped)
        int uvRowStride = planes[1].getRowStride();
        int uvPixelStride = planes[1].getPixelStride();
        int pos = ySize;

        for (int row = 0; row < height / 2; row++) {
            for (int col = 0; col < width / 2; col++) {
                int vuPos = col * uvPixelStride + row * uvRowStride;
                nv21[pos++] = vBuffer.get(vuPos);
                nv21[pos++] = uBuffer.get(vuPos);
            }
        }

        return nv21;
    }

    // Alternative method using direct YUV-to-RGB color space conversion
    private Bitmap yuvImageToBitmapDirect(Image image) {
        if (image.getFormat() != ImageFormat.YUV_420_888) {
            throw new IllegalArgumentException("Only YUV_420_888 format is supported");
        }

        int width = image.getWidth();
        int height = image.getHeight();
        Bitmap bitmap = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888);

        // Get YUV data
        Image.Plane[] planes = image.getPlanes();
        ByteBuffer yBuffer = planes[0].getBuffer();
        ByteBuffer uBuffer = planes[1].getBuffer();
        ByteBuffer vBuffer = planes[2].getBuffer();

        int yRowStride = planes[0].getRowStride();
        int uvRowStride = planes[1].getRowStride();
        int uvPixelStride = planes[1].getPixelStride();

        // Convert YUV to RGB directly, pixel by pixel
        int[] rgbData = new int[width * height];

        for (int j = 0; j < height; j++) {
            for (int i = 0; i < width; i++) {
                int yIndex = j * yRowStride + i;
                int uvJ = j / 2;
                int uvI = i / 2;
                int uvIndex = uvJ * uvRowStride + uvI * uvPixelStride;

                int y = yBuffer.get(yIndex) & 0xFF;
                int u = uBuffer.get(uvIndex) & 0xFF;
                int v = vBuffer.get(uvIndex) & 0xFF;

                // YUV to RGB conversion
                y = y - 16;
                u = u - 128;
                v = v - 128;

                int r = (int) (1.164 * y + 1.596 * v);
                int g = (int) (1.164 * y - 0.813 * v - 0.391 * u);
                int b = (int) (1.164 * y + 2.018 * u);

                // Clamp RGB values
                r = r > 255 ? 255 : (r < 0 ? 0 : r);
                g = g > 255 ? 255 : (g < 0 ? 0 : g);
                b = b > 255 ? 255 : (b < 0 ? 0 : b);

                rgbData[j * width + i] = 0xFF000000 | (r << 16) | (g << 8) | b;
            }
        }

        bitmap.setPixels(rgbData, 0, width, 0, 0, width, height);
        return bitmap;
    }

    private void saveBitmapToFile(Bitmap bitmap, String fileName) throws IOException {
//        File picturesDir = Environment.getExternalStoragePublicDirectory(Environment.DIRECTORY_PICTURES);

        File picturesDir = Environment.getExternalStoragePublicDirectory(Environment.DIRECTORY_DCIM);
        File outputFile = new File(picturesDir, "YUV_" + generateTimestamp() + "_" + fileName + ".bmp");

        FileOutputStream outputStream = new FileOutputStream(outputFile);
        bitmap.compress(Bitmap.CompressFormat.PNG, 100, outputStream);
        outputStream.flush();
        outputStream.close();
    }
}