package org.lsposed.lspatch.util;

import android.os.SharedMemory;
import android.system.ErrnoException;
import android.system.OsConstants;
import android.util.Log;

import org.lsposed.lspd.models.PreLoadedApk;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.channels.Channels;
import java.nio.channels.ReadableByteChannel;
import java.util.ArrayList;
import java.util.List;
import java.util.zip.ZipFile;

public class ModuleLoader {

    private static final String TAG = "LSPatch";

    private static void readFully(ReadableByteChannel channel, java.nio.ByteBuffer byteBuffer) throws IOException {
        while (byteBuffer.hasRemaining()) {
            if (channel.read(byteBuffer) < 0) {
                throw new IOException("Unexpected end of dex stream");
            }
        }
    }

    private static void readDexes(ZipFile apkFile, List<SharedMemory> preLoadedDexes) {
        int secondary = 2;
        for (var dexFile = apkFile.getEntry("classes.dex"); dexFile != null;
             dexFile = apkFile.getEntry("classes" + secondary + ".dex"), secondary++) {
            try (var in = apkFile.getInputStream(dexFile)) {
                var dexSize = Math.toIntExact(dexFile.getSize());
                Log.i(TAG, "ModuleLoader: preload dex entry=" + dexFile.getName()
                        + " size=" + dexFile.getSize()
                        + " compressedSize=" + dexFile.getCompressedSize()
                        + " crc=" + dexFile.getCrc()
                        + " apk=" + apkFile.getName());
                var memory = SharedMemory.create(null, dexSize);
                var byteBuffer = memory.mapReadWrite();
                try (var channel = Channels.newChannel(in)) {
                    readFully(channel, byteBuffer);
                }
                SharedMemory.unmap(byteBuffer);
                memory.setProtect(OsConstants.PROT_READ);
                preLoadedDexes.add(memory);
                Log.i(TAG, "ModuleLoader: preloaded dex entry=" + dexFile.getName()
                        + " sharedMemorySize=" + memory.getSize()
                        + " dexCount=" + preLoadedDexes.size());
            } catch (IOException | ErrnoException e) {
                Log.w(TAG, "Can not load " + dexFile + " in " + apkFile, e);
            }
        }
    }

    private static void readName(ZipFile apkFile, String initName, List<String> names) {
        var initEntry = apkFile.getEntry(initName);
        if (initEntry == null) {
            Log.i(TAG, "ModuleLoader: init entry missing " + initName + " apk=" + apkFile.getName());
            return;
        }
        Log.i(TAG, "ModuleLoader: read init entry " + initName
                + " size=" + initEntry.getSize()
                + " crc=" + initEntry.getCrc()
                + " apk=" + apkFile.getName());
        try (var in = apkFile.getInputStream(initEntry)) {
            var reader = new BufferedReader(new InputStreamReader(in));
            String name;
            while ((name = reader.readLine()) != null) {
                name = name.trim();
                if (name.isEmpty() || name.startsWith("#")) continue;
                names.add(name);
                Log.i(TAG, "ModuleLoader: init entry " + initName + " contains " + name);
            }
        } catch (IOException e) {
            Log.e(TAG, "Can not open " + initEntry, e);
        }
    }

    public static PreLoadedApk loadModule(String path) {
        Log.i(TAG, "ModuleLoader: loadModule start path=" + path);
        if (path == null) {
            Log.w(TAG, "ModuleLoader: loadModule path is null");
            return null;
        }
        var file = new PreLoadedApk();
        var preLoadedDexes = new ArrayList<SharedMemory>();
        var moduleClassNames = new ArrayList<String>(1);
        var moduleLibraryNames = new ArrayList<String>(1);
        try (var apkFile = new ZipFile(path)) {
            readDexes(apkFile, preLoadedDexes);
            readName(apkFile, "assets/xposed_init", moduleClassNames);
            readName(apkFile, "assets/native_init", moduleLibraryNames);
        } catch (IOException e) {
            Log.e(TAG, "Can not open " + path, e);
            return null;
        }
        Log.i(TAG, "ModuleLoader: loadModule parsed path=" + path
                + " dexCount=" + preLoadedDexes.size()
                + " moduleClassNames=" + moduleClassNames
                + " moduleLibraryNames=" + moduleLibraryNames);
        if (preLoadedDexes.isEmpty()) {
            Log.w(TAG, "ModuleLoader: loadModule no dexes path=" + path);
            return null;
        }
        if (moduleClassNames.isEmpty()) {
            Log.w(TAG, "ModuleLoader: loadModule no xposed classes path=" + path);
            return null;
        }
        file.preLoadedDexes = preLoadedDexes;
        file.moduleClassNames = moduleClassNames;
        file.moduleLibraryNames = moduleLibraryNames;
        Log.i(TAG, "ModuleLoader: loadModule done path=" + path);
        return file;
    }
}