/*
 *   This file is a part of the SteamWar software.
 *
 *   Copyright (C) 2026 SteamWar.de-Serverteam
 *
 *   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 de.steamwar;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.security.ProtectionDomain;
import java.util.*;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.logging.Logger;
import java.util.stream.Collectors;

/**
 * Transforms class bytecode to apply access widening rules.
 * <p>
 * Also monitors for new plugin ClassLoaders appearing (when plugins load after
 * the agent attaches) and automatically picks up their .accesswidener files.
 */
public class WideningTransformer implements ClassFileTransformer {

    private static final Logger LOG = Logger.getLogger("WidenerTransformer");

    private final Instrumentation instrumentation;

    /**
     * All entries collected across all plugins. Thread-safe for concurrent plugin loads.
     */
    private final List<AccessWidenerEntry> entries = new CopyOnWriteArrayList<>();

    /**
     * ClassLoaders we have already scanned, to avoid re-scanning.
     */
    private final Set<ClassLoader> scannedLoaders = Collections.synchronizedSet(Collections.newSetFromMap(new IdentityHashMap<>()));

    public WideningTransformer(Instrumentation instrumentation) {
        this.instrumentation = instrumentation;
    }

    /**
     * Add entries — used during initial scan before the transformer is registered.
     */
    public void addEntries(List<AccessWidenerEntry> newEntries) {
        entries.addAll(newEntries);
    }

    // -------------------------------------------------------------------------
    // ClassFileTransformer
    // -------------------------------------------------------------------------

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) {

        // Check for new plugin classloaders we haven't seen yet
        if (loader != null && AccessWidenerScanner.isPluginClassLoader(loader) && scannedLoaders.add(loader)) {
            onNewPluginClassLoader(loader);
        }

        // No entries for this class — skip transformation entirely
        if (className == null) return classfileBuffer;
        boolean relevant = entries.stream().anyMatch(e -> e.targets(className));
        if (!relevant) return classfileBuffer;

        // Apply widening via ASM
        try {
            ClassReader cr = new ClassReader(classfileBuffer);
            ClassWriter cw = new ClassWriter(cr, 0);
            cr.accept(new ClassTransformer(cw, className, entries), 0);
            return cw.toByteArray();
        } catch (Exception e) {
            LOG.warning("[AccessWidener] Failed to transform " + className + ": " + e.getMessage());
            return classfileBuffer;
        }
    }

    // -------------------------------------------------------------------------
    // New ClassLoader detection
    // -------------------------------------------------------------------------

    private void onNewPluginClassLoader(ClassLoader loader) {
        List<AccessWidenerEntry> newEntries = AccessWidenerScanner.scanClassLoader(loader);
        if (newEntries.isEmpty()) return;

        entries.addAll(newEntries);
        LOG.info("[AccessWidener] Picked up " + newEntries.size() + " new entr" + (newEntries.size() == 1 ? "y" : "ies") + " from " + loader);

        // Retransform already-loaded classes that are now targeted
        scheduleRetransform(newEntries);
    }

    /**
     * Retransformation cannot be called from within a transform() call,
     * so we dispatch it to a short-lived daemon thread.
     */
    private void scheduleRetransform(List<AccessWidenerEntry> newEntries) {
        Set<String> targets = newEntries.stream().map(e -> e.target().replace('/', '.')).collect(Collectors.toSet());

        Thread t = new Thread(() -> retransform(targets), "access-widener-retransform");
        t.setDaemon(true);
        t.start();
    }

    private void retransform(Set<String> dotNames) {
        List<Class> toRetransform = Arrays.stream(instrumentation.getAllLoadedClasses())
                .filter(c -> dotNames.contains(c.getName()))
                .filter(instrumentation::isModifiableClass)
                .collect(Collectors.toList());

        if (toRetransform.isEmpty()) return;

        LOG.info("[AccessWidener] Retransforming " + toRetransform.size() + " class(es).");
        List<String> failed = new ArrayList<>();

        for (Class<?> clazz : toRetransform) {
            try {
                instrumentation.retransformClasses(clazz);
            } catch (UnmodifiableClassException | UnsupportedOperationException e) {
                failed.add(clazz.getName());
            }
        }

        if (!failed.isEmpty()) {
            LOG.warning("[AccessWidener] The following classes were already loaded before the agent"
                    + " and cannot be retransformed. Add -javaagent: to your start script"
                    + " to fix this:");
            failed.forEach(name -> LOG.warning("  - " + name));
        }
    }
}