package com.bonus.sdk; // 

import java.awt.image.BufferedImage;
import java.awt.image.DataBufferByte;
import java.io.*;
import java.nio.file.Files;
import java.nio.file.StandardCopyOption;
import java.util.HashSet;
import java.util.Set;

public class YoloSdk implements AutoCloseable {

    private long nativeHandle; 
    private static final Set<String> loadedLibraries = new HashSet<>();

    static {
        try {
            loadSdkLibrary();
        } catch (IOException e) {
            throw new RuntimeException("CRITICAL: Failed to load native YOLO SDK libraries", e);
        }
    }

    private static void loadSdkLibrary() throws IOException {
        String osName = System.getProperty("os.name").toLowerCase();
        String osArch = System.getProperty("os.arch").toLowerCase();
        String libPathInJar;
        String sdkLibName;
        String[] dependencyLibs = {}; 

        if (osName.contains("win") && osArch.contains("64")) {
 
            libPathInJar = "/lib/win-x64/";

            dependencyLibs = new String[]{
                "onnxruntime.dll",      
                "abseil_dll.dll",       
                "libprotobuf.dll",      
                "zlib1.dll",            
                "opencv_core4.dll",     
                "opencv_imgproc4.dll",  
                "opencv_dnn4.dll"       
                
            };
            
            sdkLibName = "my_yolo_sdk.dll";

        } else if ((osName.contains("nix") || osName.contains("nux")) && osArch.contains("64")) {
  
            libPathInJar = "/lib/linux-x86_64/";
            
            dependencyLibs = new String[]{
                "libonnxruntime.so.1.23.2", 
                "libopencv_core.so.4.6.0",    
                "libopencv_imgproc.so.4.6.0", 
                "libopencv_dnn.so.4.6.0"      
            };

            sdkLibName = "libmy_yolo_sdk.so";

        } else {
            throw new UnsupportedOperationException("Unsupported OS/Arch: " + osName + "/" + osArch);
        }

        for (String lib : dependencyLibs) {
            loadLibraryFromJar(libPathInJar + lib);
        }

        loadLibraryFromJar(libPathInJar + sdkLibName);
    }

 
    private static void loadLibraryFromJar(String path) throws IOException {
        String libName = new File(path).getName();
        if (loadedLibraries.contains(libName)) {
            return; 
        }

        try (InputStream in = YoloSdk.class.getResourceAsStream(path)) {
            if (in == null) {
                throw new FileNotFoundException("Library " + path + " not found in JAR.");
            }

            File tempFile = File.createTempFile(libName, ".tmp");
            tempFile.deleteOnExit(); // 

            Files.copy(in, tempFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
            System.load(tempFile.getAbsolutePath());
            loadedLibraries.add(libName);
        }
    }

    private native long nativeInit(String modelPath, int inputWidth, int inputHeight);
    private native void nativeRelease(long handle);
    private native Detection[] nativePredict(
            long handle, byte[] bgrBytes, int imageWidth, int imageHeight,
            float confThreshold, float iouThreshold
    );

    // --- 
    public YoloSdk(String modelPath, int inputWidth, int inputHeight) {
        this.nativeHandle = nativeInit(modelPath, inputWidth, inputHeight);
        if (this.nativeHandle == 0) {
            throw new RuntimeException("Failed to initialize native YOLO SDK. Check logs.");
        }
    }

    public Detection[] predict(BufferedImage image, float confThreshold, float iouThreshold) {
        if (this.nativeHandle == 0) {
            throw new IllegalStateException("SDK already closed or failed to initialize.");
        }
        byte[] bgrBytes = getBgrBytes(image);
        return nativePredict(
                this.nativeHandle, bgrBytes, image.getWidth(), image.getHeight(),
                confThreshold, iouThreshold
        );
    }

    @Override
    public void close() {
        if (this.nativeHandle != 0) {
            nativeRelease(this.nativeHandle);
            this.nativeHandle = 0;
        }
    }

    private byte[] getBgrBytes(BufferedImage image) {
        if (image.getType() == BufferedImage.TYPE_3BYTE_BGR) {
            return ((DataBufferByte) image.getRaster().getDataBuffer()).getData();
        }
        BufferedImage bgrImage = new BufferedImage(image.getWidth(), image.getHeight(), BufferedImage.TYPE_3BYTE_BGR);
        bgrImage.getGraphics().drawImage(image, 0, 0, null);
        return ((DataBufferByte) bgrImage.getRaster().getDataBuffer()).getData();
    }
}