package org.lsposed.lspd.util;

import static de.robv.android.xposed.XposedBridge.TAG;

import android.os.Build;
import android.os.SharedMemory;
import android.system.ErrnoException;
import android.system.Os;
import android.system.OsConstants;

import androidx.annotation.NonNull;
import androidx.annotation.RequiresApi;

import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Enumeration;
import java.util.List;
import java.util.Objects;
import java.util.jar.JarFile;
import java.util.zip.ZipEntry;

import org.lsposed.lspd.util.Utils.Log;

import hidden.ByteBufferDexClassLoader;
import sun.misc.CompoundEnumeration;

@SuppressWarnings("ConstantConditions")
public final class LspModuleClassLoader extends ByteBufferDexClassLoader {
    private static final String zipSeparator = "!/";
    private static final List<File> systemNativeLibraryDirs =
            splitPaths(System.getProperty("java.library.path"));
    private final String apk;
    private final List<File> nativeLibraryDirs = new ArrayList<>();

    private static List<File> splitPaths(String searchPath) {
        var result = new ArrayList<File>();
        if (searchPath == null) return result;
        for (var path : searchPath.split(File.pathSeparator)) {
            result.add(new File(path));
        }
        return result;
    }

    private LspModuleClassLoader(ByteBuffer[] dexBuffers,
                                 ClassLoader parent,
                                 String apk) {
        super(dexBuffers, parent);
        this.apk = apk;
    }

    @RequiresApi(Build.VERSION_CODES.Q)
    private LspModuleClassLoader(ByteBuffer[] dexBuffers,
                                 String librarySearchPath,
                                 ClassLoader parent,
                                 String apk) {
        super(dexBuffers, librarySearchPath, parent);
        initNativeLibraryDirs(librarySearchPath);
        this.apk = apk;
    }

    private void initNativeLibraryDirs(String librarySearchPath) {
        nativeLibraryDirs.addAll(splitPaths(librarySearchPath));
        nativeLibraryDirs.addAll(systemNativeLibraryDirs);
    }

    private static boolean shouldTraceClass(String name) {
        return name.startsWith("de.robv.android.xposed.")
                || name.startsWith("org.lsposed.")
                || name.startsWith("me.eternal.")
                || name.startsWith("com.snapchat.")
                || name.equals("android.content.res.XResources")
                || name.equals("xposed.dummy.XResourcesSuperClass")
                || name.contains("XResources")
                || name.contains("Purrfect");
    }

    @Override
    protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
        var trace = shouldTraceClass(name);
        if (trace) {
            Log.i(TAG, "LspModuleClassLoader: loadClass start name=" + name
                    + " resolve=" + resolve
                    + " loader=" + this
                    + " parent=" + getParent());
        }
        var cl = findLoadedClass(name);
        if (cl != null) {
            if (trace) {
                Log.i(TAG, "LspModuleClassLoader: loadClass hit loaded name=" + name
                        + " classLoader=" + cl.getClassLoader());
            }
            return cl;
        }
        try {
            cl = Object.class.getClassLoader().loadClass(name);
            if (trace) {
                Log.i(TAG, "LspModuleClassLoader: loadClass boot success name=" + name
                        + " classLoader=" + cl.getClassLoader());
            }
            return cl;
        } catch (ClassNotFoundException | NoClassDefFoundError e) {
            if (trace) {
                Log.i(TAG, "LspModuleClassLoader: loadClass boot miss name=" + name + " error=" + e);
            }
        }
        ClassNotFoundException fromSuper;
        try {
            cl = findClass(name);
            if (trace) {
                Log.i(TAG, "LspModuleClassLoader: loadClass findClass success name=" + name
                        + " classLoader=" + cl.getClassLoader());
            }
            return cl;
        } catch (ClassNotFoundException ex) {
            fromSuper = ex;
            if (trace) {
                Log.i(TAG, "LspModuleClassLoader: loadClass findClass miss name=" + name + " error=" + ex);
            }
        }
        try {
            cl = getParent().loadClass(name);
            if (trace) {
                Log.i(TAG, "LspModuleClassLoader: loadClass parent success name=" + name
                        + " classLoader=" + cl.getClassLoader());
            }
            return cl;
        } catch (ClassNotFoundException cnfe) {
            if (trace) {
                Log.e(TAG, "LspModuleClassLoader: loadClass failed name=" + name, fromSuper);
            }
            throw fromSuper;
        }
    }

    @Override
    public String findLibrary(String libraryName) {
        var fileName = System.mapLibraryName(libraryName);
        Log.i(TAG, "LspModuleClassLoader: findLibrary start libraryName=" + libraryName
                + " fileName=" + fileName
                + " dirs=" + nativeLibraryDirs);
        for (var file : nativeLibraryDirs) {
            var path = file.getPath();
            if (path.contains(zipSeparator)) {
                var split = path.split(zipSeparator, 2);
                try (var jarFile = new JarFile(split[0])) {
                    var entryName = split[1] + '/' + fileName;
                    var entry = jarFile.getEntry(entryName);
                    if (entry != null && entry.getMethod() == ZipEntry.STORED) {
                        Log.i(TAG, "LspModuleClassLoader: findLibrary found stored zip entry " + entryName
                                + " in " + split[0]);
                        return split[0] + zipSeparator + entryName;
                    }
                    Log.i(TAG, "LspModuleClassLoader: findLibrary zip miss entry=" + entryName
                            + " jar=" + split[0]
                            + " entry=" + entry);
                } catch (IOException e) {
                    Log.e(TAG, "Can not open " + split[0], e);
                }
            } else if (file.isDirectory()) {
                var entryPath = new File(file, fileName).getPath();
                try {
                    var fd = Os.open(entryPath, OsConstants.O_RDONLY, 0);
                    Os.close(fd);
                    Log.i(TAG, "LspModuleClassLoader: findLibrary found file " + entryPath);
                    return entryPath;
                } catch (ErrnoException ignored) {
                }
            }
        }
        Log.i(TAG, "LspModuleClassLoader: findLibrary miss libraryName=" + libraryName);
        return null;
    }

    @Override
    public String getLdLibraryPath() {
        var result = new StringBuilder();
        for (var directory : nativeLibraryDirs) {
            if (result.length() > 0) {
                result.append(':');
            }
            result.append(directory);
        }
        return result.toString();
    }

    @Override
    protected URL findResource(String name) {
        if (name.startsWith("assets/") || name.contains("xposed") || name.contains("lspatch")) {
            Log.i(TAG, "LspModuleClassLoader: findResource start name=" + name + " apk=" + apk);
        }
        try {
            var urlHandler = new ClassPathURLStreamHandler(apk);
            var url = urlHandler.getEntryUrlOrNull(name);
            if (url == null) {
                // noinspection FinalizeCalledExplicitly
                urlHandler.finalize();
            }
            if (name.startsWith("assets/") || name.contains("xposed") || name.contains("lspatch")) {
                Log.i(TAG, "LspModuleClassLoader: findResource result name=" + name + " url=" + url);
            }
            return url;
        } catch (IOException e) {
            if (name.startsWith("assets/") || name.contains("xposed") || name.contains("lspatch")) {
                Log.e(TAG, "LspModuleClassLoader: findResource failed name=" + name, e);
            }
            return null;
        }
    }

    @Override
    protected Enumeration<URL> findResources(String name) {
        var result = new ArrayList<URL>();
        var url = findResource(name);
        if (url != null) result.add(url);
        return Collections.enumeration(result);
    }

    @Override
    public URL getResource(String name) {
        var resource = Object.class.getClassLoader().getResource(name);
        if (resource != null) return resource;
        resource = findResource(name);
        if (resource != null) return resource;
        final var cl = getParent();
        return (cl == null) ? null : cl.getResource(name);
    }

    @Override
    public Enumeration<URL> getResources(String name) throws IOException {
        @SuppressWarnings("unchecked") final var resources = (Enumeration<URL>[]) new Enumeration<?>[]{
                Object.class.getClassLoader().getResources(name),
                findResources(name),
                getParent() == null ? null : getParent().getResources(name)};
        return new CompoundEnumeration<>(resources);
    }

    @NonNull
    @Override
    public String toString() {
        if (apk == null) return "LspModuleClassLoader[instantiating]";
        return "LspModuleClassLoader[module=" + apk + ", " + super.toString() + "]";
    }

    public static ClassLoader loadApk(String apk,
                                      List<SharedMemory> dexes,
                                      String librarySearchPath,
                                      ClassLoader parent) {
        Log.i(TAG, "LspModuleClassLoader: loadApk start apk=" + apk
                + " dexCount=" + (dexes == null ? -1 : dexes.size())
                + " librarySearchPath=" + librarySearchPath
                + " parent=" + parent
                + " sdk=" + Build.VERSION.SDK_INT);
        var dexBuffers = dexes.stream().parallel().map(dex -> {
            try {
                var buffer = dex.mapReadOnly();
                Log.i(TAG, "LspModuleClassLoader: mapped dex sharedMemory=" + dex
                        + " bufferCapacity=" + buffer.capacity()
                        + " remaining=" + buffer.remaining());
                return buffer;
            } catch (ErrnoException e) {
                Log.w(TAG, "Can not map " + dex, e);
                return null;
            }
        }).filter(Objects::nonNull).toArray(ByteBuffer[]::new);
        Log.i(TAG, "LspModuleClassLoader: loadApk mapped dexBuffers=" + dexBuffers.length);
        LspModuleClassLoader cl;
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.Q) {
            cl = new LspModuleClassLoader(dexBuffers, librarySearchPath, parent, apk);
        } else {
            cl = new LspModuleClassLoader(dexBuffers, parent, apk);
            cl.initNativeLibraryDirs(librarySearchPath);
        }
        Log.i(TAG, "LspModuleClassLoader: loadApk created classLoader=" + cl
                + " nativeLibraryDirs=" + cl.nativeLibraryDirs);
        Arrays.stream(dexBuffers).parallel().forEach(SharedMemory::unmap);
        dexes.stream().parallel().forEach(SharedMemory::close);
        Log.i(TAG, "LspModuleClassLoader: loadApk unmapped buffers and closed shared memory");
        return cl;
    }
}