/*
 * SoulFire
 * Copyright (C) 2026  AlexProgrammerDE
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
package com.soulfiremc.launcher;

import com.google.common.reflect.ClassPath;
import net.lenni0451.classtransform.utils.tree.IClassProvider;
import org.checkerframework.checker.nullness.qual.NonNull;
import org.checkerframework.checker.nullness.qual.Nullable;

import javax.annotation.Nonnull;
import java.io.ByteArrayOutputStream;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.function.Supplier;

import static net.lenni0451.classtransform.utils.ASMUtils.slash;
import static net.lenni0451.classtransform.utils.Sneaky.sneakySupply;

public final class CustomClassProvider implements IClassProvider {
  private final ClassPath[] classPaths;
  private final ClassLoader[] classLoaders;

  public CustomClassProvider(List<ClassLoader> classLoaders) {
    this.classLoaders = classLoaders.toArray(new ClassLoader[0]);

    try {
      this.classPaths = new ClassPath[classLoaders.size()];
      for (var i = 0; i < classLoaders.size(); i++) {
        this.classPaths[i] = ClassPath.from(classLoaders.get(i));
      }
    } catch (Throwable t) {
      throw new RuntimeException("Failed to initialize ClassPath", t);
    }
  }

  @Override
  public byte @NonNull [] getClass(@NonNull String name) throws ClassNotFoundException {
    for (var classLoader : this.classLoaders) {
      var bytes = this.getClassFromLoader(classLoader, name);
      if (bytes != null) {
        return bytes;
      }
    }

    throw new ClassNotFoundException("Class not found: " + name);
  }

  private byte @Nullable [] getClassFromLoader(ClassLoader classLoader, String name) {
    try (var is = classLoader.getResourceAsStream(slash(name) + ".class")) {
      Objects.requireNonNull(is, "Class input stream is null");
      var baos = new ByteArrayOutputStream();
      var buf = new byte[1024];
      int len;
      while ((len = is.read(buf)) > 0) {
        baos.write(buf, 0, len);
      }
      return baos.toByteArray();
    } catch (Throwable _) {
      return null;
    }
  }

  @Override
  @Nonnull
  public Map<String, Supplier<byte[]>> getAllClasses() {
    var map = new HashMap<String, Supplier<byte[]>>();
    for (var classPath : this.classPaths) {
      for (var classInfo : classPath.getAllClasses()) {
        map.put(classInfo.getName(), sneakySupply(() -> this.getClass(classInfo.getName())));
      }
    }
    return map;
  }
}
