/* MIT. A typed value encoder for the hash-pinned TLC build, not a TLA evaluator. */
import java.nio.file.*;
import java.nio.channels.FileChannel;
import java.nio.charset.StandardCharsets;
import java.util.*;
import com.google.gson.*;
import tlc2.value.impl.*;

public final class FormalWikiTrace {
    private static final Gson JSON = new GsonBuilder().create();
    private int remaining = Integer.getInteger("formalwiki.maxnodes", 500000);
    private final int maxDepth = Integer.getInteger("formalwiki.maxdepth", 128);

    @tlc2.overrides.TLAPlusOperator(identifier="TraceEmit", module="FormalWikiTrace", warn=true)
    public static Value TraceEmit(Value value) throws Exception {
        String run = System.getProperty("formalwiki.run");
        String input = System.getProperty("formalwiki.input");
        if (run == null || input == null || !input.matches("[0-9a-f]{64}")) {
            throw new IllegalArgumentException("Capture requires run and exact input identity");
        }
        Path target = Path.of(System.getProperty("formalwiki.capture"));
        if (target.isAbsolute() || target.normalize().startsWith("..")) {
            throw new IllegalArgumentException("Capture path must remain in the run directory");
        }
        target = target.toAbsolutePath().normalize();
        for (Path parent = target; parent != null; parent = parent.getParent()) {
            if (Files.isSymbolicLink(parent)) throw new IllegalArgumentException("Symlink capture path");
        }
        if (Files.exists(target)) throw new FileAlreadyExistsException(target.toString());
        FormalWikiTrace encoder = new FormalWikiTrace();
        JsonObject capture = new JsonObject();
        capture.addProperty("schema_version", 1);
        capture.addProperty("kind", "typed-tlc-capture");
        capture.addProperty("run_id", run);
        capture.addProperty("input_manifest_sha256", input);
        capture.add("counterexample", encoder.encode(value, 0));
        JsonArray atoms = new JsonArray();
        if (ModelValue.mvs != null) for (ModelValue atom : ModelValue.mvs) {
            JsonObject row = new JsonObject();
            row.addProperty("index", atom.index);
            row.addProperty("name", atom.val.toString());
            atoms.add(row);
        }
        capture.add("model_values", atoms);
        byte[] bytes = (JSON.toJson(capture) + "\n").getBytes(StandardCharsets.UTF_8);
        if (bytes.length > Integer.getInteger("formalwiki.maxbytes", 16777216)) {
            throw new IllegalArgumentException("Typed capture exceeds byte bound");
        }
        Path temporary = Files.createTempFile(target.getParent(), ".capture-", ".partial");
        try {
            Files.write(temporary, bytes, StandardOpenOption.TRUNCATE_EXISTING);
            try (FileChannel file = FileChannel.open(temporary, StandardOpenOption.WRITE)) {
                file.force(true);
            }
            // Exclusive link publication refuses overwrite and exposes only complete bytes.
            Files.createLink(target, temporary);
        } finally {
            Files.deleteIfExists(temporary);
        }
        return BoolValue.ValTrue;
    }

    private static JsonObject tagged(String kind) {
        JsonObject result = new JsonObject();
        result.addProperty("kind", kind);
        return result;
    }

    private JsonElement encode(Value value, int depth) {
        if (--remaining < 0 || depth > maxDepth) {
            throw new IllegalArgumentException("Typed capture exceeds node/depth bound");
        }
        if (value instanceof IntValue) {
            JsonObject result = tagged("integer");
            result.addProperty("value", Integer.toString(((IntValue)value).val));
            return result;
        }
        if (value instanceof BoolValue) {
            JsonObject result = tagged("boolean");
            result.addProperty("value", ((BoolValue)value).val);
            return result;
        }
        if (value instanceof StringValue) {
            JsonObject result = tagged("string");
            result.addProperty("value", ((StringValue)value).val.toString());
            return result;
        }
        if (value instanceof ModelValue) {
            JsonObject result = tagged("model_value");
            result.addProperty("value", ((ModelValue)value).val.toString());
            return result;
        }
        if (value instanceof TupleValue) {
            JsonObject result = tagged("tuple");
            JsonArray items = new JsonArray();
            for (Value element : ((TupleValue)value).elems) items.add(encode(element, depth + 1));
            result.add("items", items);
            return result;
        }
        if (value instanceof RecordValue) {
            RecordValue record = (RecordValue)value;
            JsonObject result = tagged("record");
            TreeMap<String, JsonElement> sorted = new TreeMap<>();
            for (int index = 0; index < record.names.length; index++) {
                sorted.put(record.names[index].toString(), encode(record.values[index], depth + 1));
            }
            JsonObject fields = new JsonObject();
            sorted.forEach(fields::add);
            result.add("fields", fields);
            return result;
        }
        if (value instanceof SetEnumValue) {
            JsonObject result = tagged("set");
            ArrayList<JsonElement> sorted = new ArrayList<>();
            for (Value element : ((SetEnumValue)value).elems.toArray()) sorted.add(encode(element, depth + 1));
            sorted.sort(Comparator.comparing(JSON::toJson));
            JsonArray items = new JsonArray();
            sorted.forEach(items::add);
            result.add("items", items);
            return result;
        }
        if (value instanceof FcnRcdValue) {
            FcnRcdValue function = (FcnRcdValue)value;
            JsonObject result = tagged("function");
            Value[] domain = function.getDomainAsValues();
            ArrayList<JsonObject> sorted = new ArrayList<>();
            for (int index = 0; index < domain.length; index++) {
                JsonObject row = new JsonObject();
                row.add("key", encode(domain[index], depth + 1));
                row.add("value", encode(function.values[index], depth + 1));
                sorted.add(row);
            }
            sorted.sort(Comparator.comparing(row -> JSON.toJson(row.get("key"))));
            JsonArray entries = new JsonArray();
            sorted.forEach(entries::add);
            result.add("entries", entries);
            return result;
        }
        throw new IllegalArgumentException("Unsupported typed capture value: " + value.getClass().getName());
    }
}