package jmv;

import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.LinkedHashMap;
import java.util.Map;

/**
 * Entry point run inside CheerpJ (or natively for tests).
 *
 * One-shot: {@code sourcePath [stdinPath|-] [outputPath] [apiJar] [stepBudget] [warm]}.
 * Resident: {@code --serve [apiJar]} warms the compiler once, signals {@link #ready}, then loops on
 * {@link #nextJob} (a JS-implemented async native that resolves when the page submits a job).
 * Each job reads /str/Main.java and /str/stdin.txt and answers through {@link #emit}.
 * Keeping ECJ loaded turns the ~16 s cold compile into a ~0.2 s warm compile per run.
 */
public final class Visualizer {
    public static native void emit(String json);
    public static native void ready(long warmupNs);
    /** Returns "stepBudget|cold" or "stepBudget|warm"; empty string ends the loop. */
    public static native String nextJob();

    static final String WARMUP = "import java.util.*;\n"
        + "public class Warmup { interface S { double a(); } record P(int x) {} enum C { R, G }\n"
        + "  public static void main(String[] args) { List<Integer> l = new ArrayList<>(); l.add(1);\n"
        + "    Map<String, P> m = new HashMap<>(); m.put(\"k\", new P(2)); Runnable r = () -> l.add(3); r.run();\n"
        + "    try { System.out.println(l + \" \" + m + C.R + args.length); } catch (RuntimeException e) { throw e; } } }\n";

    static final class MemoryLoader extends ClassLoader {
        final Map<String, byte[]> classes;
        MemoryLoader(Map<String, byte[]> classes) { super(Visualizer.class.getClassLoader()); this.classes = classes; }
        @Override protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
            if (!classes.containsKey(name)) return super.loadClass(name, resolve);
            Class<?> c = findLoadedClass(name);
            if (c == null) { byte[] b = classes.get(name); c = defineClass(name, b, 0, b.length); }
            if (resolve) resolveClass(c);
            return c;
        }
    }

    public static void main(String[] args) throws Exception {
        if (args.length > 0 && args[0].equals("--serve")) { serve(args.length > 1 ? args[1] : "/app/web/vendor/java17-api.jar"); return; }
        String source = Files.readString(Paths.get(args[0]));
        String stdin = args.length > 1 && !args[1].equals("-") ? Files.readString(Paths.get(args[1])) : "";
        String outputPath = args.length > 2 && !args[2].equals("-") ? args[2] : null;
        String apiJar = args.length > 3 ? args[3] : "/app/web/vendor/java17-api.jar";
        int stepBudget = args.length > 4 ? Integer.parseInt(args[4]) : 1000;
        boolean warm = args.length > 5 && args[5].equals("warm");
        boolean objects = args.length > 6 && args[6].contains("objects");
        String json = run(source, stdin, apiJar, stepBudget, warm, objects);
        if (outputPath != null) Files.writeString(Paths.get(outputPath), json); else emit(json);
    }

    static void serve(String apiJar) {
        long t0 = System.nanoTime();
        try { Compiler.compile(WARMUP, apiJar); } catch (Throwable e) { System.err.println("warmup failed: " + e); }
        ready(System.nanoTime() - t0);
        while (true) {
            String job = nextJob();
            if (job == null || job.isEmpty()) return;
            String[] parts = job.split("\\|");
            int budget = 1000;
            try { budget = Integer.parseInt(parts[0]); } catch (NumberFormatException ignored) { }
            boolean warm = parts.length > 1 && parts[1].equals("warm");
            boolean objects = parts.length > 2 && parts[2].contains("objects");
            String json;
            try {
                String source = Files.readString(Paths.get("/str/Main.java"));
                String stdin = Files.readString(Paths.get("/str/stdin.txt"));
                json = run(source, stdin, apiJar, budget, warm, objects);
            } catch (Throwable e) {
                json = "{\"status\":\"internal_error\",\"message\":" + Tracer.json(e.toString()) + ",\"problems\":[],\"trace\":[]}";
            }
            emit(json);
        }
    }

    public static String run(String source, String stdin, String apiJar, int stepBudget) { return run(source, stdin, apiJar, stepBudget, false, false); }

    /**
     * warm=true compiles twice and reports the second (class-loading-free) compile time as compileWarmNs.
     * objects=true records String and boxed primitives as heap objects so their identity is visible.
     */
    public static String run(String source, String stdin, String apiJar, int stepBudget, boolean warm, boolean objects) {
        StringBuilder out = new StringBuilder(4096);
        out.append("{\"code\":").append(Tracer.json(source)).append(",\"stdin\":").append(Tracer.json(stdin));
        long t0 = System.nanoTime();
        Compiler.Result compiled;
        try { compiled = Compiler.compile(source, apiJar); }
        catch (Throwable e) {
            return out.append(",\"status\":\"internal_error\",\"message\":").append(Tracer.json("compiler failure: " + e)).append(",\"problems\":[],\"trace\":[]}").toString();
        }
        long compileNs = System.nanoTime() - t0;
        long compileWarmNs = -1;
        if (warm) { t0 = System.nanoTime(); Compiler.compile(source, apiJar); compileWarmNs = System.nanoTime() - t0; }
        out.append(",\"problems\":[");
        for (int i = 0; i < compiled.problems.size(); i++) {
            Compiler.Problem p = compiled.problems.get(i);
            if (i > 0) out.append(',');
            out.append("{\"line\":").append(p.line).append(",\"error\":").append(p.error).append(",\"message\":").append(Tracer.json(p.message)).append('}');
        }
        out.append(']');
        if (compiled.hasErrors()) return out.append(",\"status\":\"compile_error\",\"trace\":[]}").toString();

        t0 = System.nanoTime();
        Instrumenter instrumenter = new Instrumenter(compiled.classes);
        Map<String, byte[]> transformed = new LinkedHashMap<>();
        try { for (Map.Entry<String, byte[]> e : compiled.classes.entrySet()) transformed.put(e.getKey(), instrumenter.transform(e.getValue())); }
        catch (Throwable e) {
            return out.append(",\"status\":\"internal_error\",\"message\":").append(Tracer.json("instrumentation failure: " + e)).append(",\"trace\":[]}").toString();
        }
        long transformNs = System.nanoTime() - t0;

        MemoryLoader loader = new MemoryLoader(transformed);
        Method main = null; String mainClass = null;
        try {
            String preferred = compiled.unitName.substring(0, compiled.unitName.length() - 5);
            for (String name : orderedCandidates(compiled, preferred)) {
                Class<?> c = loader.loadClass(name);
                try {
                    Method m = c.getDeclaredMethod("main", String[].class);
                    if (Modifier.isStatic(m.getModifiers()) && Modifier.isPublic(m.getModifiers())) { main = m; mainClass = name; break; }
                } catch (NoSuchMethodException ignored) { }
            }
        } catch (Throwable e) {
            return out.append(",\"status\":\"internal_error\",\"message\":").append(Tracer.json("class loading failure: " + e)).append(",\"trace\":[]}").toString();
        }
        if (main == null) return out.append(",\"status\":\"no_main\",\"message\":\"public static void main(String[] args) not found\",\"trace\":[]}").toString();

        Tracer.reset(compiled.classes.keySet(), stepBudget, objects);
        InputStream originalIn = System.in;
        System.setIn(new ByteArrayInputStream(stdin.getBytes(StandardCharsets.UTF_8)));
        Tracer.captureOutput();
        String status = "ok", message = null;
        t0 = System.nanoTime();
        try {
            main.invoke(null, (Object) new String[0]);
        } catch (InvocationTargetException e) {
            Throwable cause = e.getCause();
            if (cause instanceof Tracer.StepLimit) status = "step_limit";
            else { status = "uncaught_exception"; message = cause.toString(); }
        } catch (Throwable e) {
            status = "internal_error"; message = e.toString();
        } finally {
            Tracer.enabled = false;
            Tracer.restoreOutput();
            System.setIn(originalIn);
        }
        long runNs = System.nanoTime() - t0;
        if (Tracer.limitHit()) status = "step_limit";

        out.append(",\"status\":").append(Tracer.json(status));
        if (message != null) out.append(",\"message\":").append(Tracer.json(message));
        out.append(",\"mainClass\":").append(Tracer.json(mainClass));
        out.append(",\"stepBudget\":").append(stepBudget).append(",\"heapTruncated\":").append(Tracer.heapTruncated());
        out.append(",\"options\":{\"objects\":").append(objects).append('}');
        out.append(",\"timing\":{\"compileNs\":").append(compileNs).append(",\"compileWarmNs\":").append(compileWarmNs).append(",\"transformNs\":").append(transformNs)
           .append(",\"runNs\":").append(runNs).append(",\"snapshotNs\":").append(Tracer.snapshotNs()).append('}');
        out.append(",\"retainedObjects\":").append(Tracer.retainedObjects());
        out.append(",\"trace\":[");
        boolean first = true;
        for (String step : Tracer.steps()) { if (!first) out.append(','); first = false; out.append(step); }
        return out.append("]}").toString();
    }

    static Iterable<String> orderedCandidates(Compiler.Result compiled, String preferred) {
        java.util.ArrayList<String> names = new java.util.ArrayList<>();
        if (compiled.classes.containsKey(preferred)) names.add(preferred);
        for (String n : compiled.classes.keySet()) if (!names.contains(n) && n.indexOf('$') < 0) names.add(n);
        for (String n : compiled.classes.keySet()) if (!names.contains(n)) names.add(n);
        return names;
    }
}