/* MIT. Lossless typed DAG capture for the exact hash-pinned TLC build. */
import java.nio.file.*;
import java.nio.channels.FileChannel;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.util.*;
import com.google.gson.*;
import tlc2.value.impl.*;
import util.UniqueString;

public final class FormalWikiTrace {
    private static final Gson JSON = new GsonBuilder().create();
    private final IdentityHashMap<Value,String> fingerprints = new IdentityHashMap<>();
    private final Set<Value> active = Collections.newSetFromMap(new IdentityHashMap<>());
    private final HashMap<String,String> fingerprintShapes = new HashMap<>();
    private final HashMap<String,Integer> heights = new HashMap<>();
    private final IdentityHashMap<FcnRcdValue,Value[]> functionDomains = new IdentityHashMap<>();
    private final IdentityHashMap<Value,Integer> identities = new IdentityHashMap<>();
    private final HashMap<String,Integer> structures = new HashMap<>();
    private final ArrayList<JsonObject> pool = new ArrayList<>();
    private final HashMap<String,Value> atoms = new HashMap<>();
    private final int maxNodes = Integer.getInteger("formalwiki.maxnodes",500000);
    private final int maxDepth = Integer.getInteger("formalwiki.maxdepth",128);
    private long edges = 0;

    private static JsonObject tagged(String kind) {
        JsonObject result = new JsonObject(); result.addProperty("kind",kind); return result;
    }
    private static JsonArray array(List<?> values) { return JSON.toJsonTree(values).getAsJsonArray(); }
    private String fingerprint(Value value,int depth) throws Exception {
        if(depth>maxDepth)throw new IllegalArgumentException("Typed capture exceeds depth bound");
        String prior=fingerprints.get(value);if(prior!=null){
            if(depth+heights.get(prior)>maxDepth)throw new IllegalArgumentException("Typed capture exceeds depth bound");
            return prior;
        }
        if(depth>maxDepth||fingerprints.size()+active.size()>=maxNodes||!active.add(value))
            throw new IllegalArgumentException("Typed capture exceeds unique-node/depth bound or contains a cycle");
        JsonObject node=describe(value,depth,true);
        String shape=JSON.toJson(node);
        String digest=HexFormat.of().formatHex(MessageDigest.getInstance("SHA-256").digest(shape.getBytes(StandardCharsets.UTF_8)));
        String old=fingerprintShapes.putIfAbsent(digest,shape);
        if(old!=null&&!old.equals(shape))throw new IllegalArgumentException("Typed capture structural digest collision");
        int height=0;
        if(node.has("items"))for(JsonElement ref:node.getAsJsonArray("items"))height=Math.max(height,heights.get(ref.getAsString())+1);
        if(node.has("fields"))for(var entry:node.getAsJsonObject("fields").entrySet())height=Math.max(height,heights.get(entry.getValue().getAsString())+1);
        if(node.has("entries"))for(JsonElement row:node.getAsJsonArray("entries"))for(String key:List.of("key","value"))height=Math.max(height,heights.get(row.getAsJsonObject().get(key).getAsString())+1);
        if(depth+height>maxDepth)throw new IllegalArgumentException("Typed capture exceeds depth bound");
        heights.put(digest,height);
        fingerprints.put(value,digest);active.remove(value);return digest;
    }
    private Object child(Value value,int depth,boolean hashing) throws Exception {
        return hashing?fingerprint(value,depth+1):visit(value,depth+1);
    }
    private JsonObject describe(Value value,int depth,boolean hashing) throws Exception {
        JsonObject node;
        if(value instanceof IntValue){node=tagged("integer");node.addProperty("value",Integer.toString(((IntValue)value).val));}
        else if(value instanceof BoolValue){node=tagged("boolean");node.addProperty("value",((BoolValue)value).val);}
        else if(value instanceof StringValue){node=tagged("string");node.addProperty("value",((StringValue)value).val.toString());}
        else if(value instanceof ModelValue){String name=((ModelValue)value).val.toString();atoms.put(name,value);node=tagged("model_value");node.addProperty("value",name);}
        else if(value instanceof TupleValue){
            node=tagged("tuple");ArrayList<Object> items=new ArrayList<>();
            for(Value v:((TupleValue)value).elems)items.add(child(v,depth,hashing));node.add("items",array(items));
        }else if(value instanceof RecordValue){
            node=tagged("record");RecordValue record=(RecordValue)value;TreeMap<String,Value> sorted=new TreeMap<>();
            for(int i=0;i<record.names.length;i++)if(sorted.put(record.names[i].toString(),record.values[i])!=null)
                throw new IllegalArgumentException("Duplicate typed record field");
            JsonObject fields=new JsonObject();for(var entry:sorted.entrySet())fields.add(entry.getKey(),JSON.toJsonTree(child(entry.getValue(),depth,hashing)));
            node.add("fields",fields);
        }else if(value instanceof SetEnumValue){
            node=tagged("set");ArrayList<Value> sorted=new ArrayList<>(Arrays.asList(((SetEnumValue)value).elems.toArray()));
            // Stable representation fingerprints determine traversal; no expanded values or run-local IDs enter this order.
            ArrayList<Object> items=new ArrayList<>();
            if(hashing){for(Value v:sorted)fingerprint(v,depth+1);sorted.sort(Comparator.comparing(v->fingerprints.get(v)));}
            else sorted.sort(Comparator.comparing(v->fingerprints.get(v)));
            for(Value v:sorted)items.add(child(v,depth,hashing));
            if(!hashing)items.sort(Comparator.comparingInt(v->(Integer)v));
            node.add("items",array(items));
        }else if(value instanceof FcnRcdValue){
            node=tagged("function");FcnRcdValue function=(FcnRcdValue)value;
            Value[] domain=functionDomains.computeIfAbsent(function,f->f.getDomainAsValues());
            ArrayList<Integer> order=new ArrayList<>();for(int i=0;i<domain.length;i++){order.add(i);if(hashing)fingerprint(domain[i],depth+1);}
            order.sort(Comparator.comparing(i->fingerprints.get(domain[i])));
            ArrayList<JsonObject> entries=new ArrayList<>();
            for(int i:order){JsonObject row=new JsonObject();row.add("key",JSON.toJsonTree(child(domain[i],depth,hashing)));row.add("value",JSON.toJsonTree(child(function.values[i],depth,hashing)));entries.add(row);}
            if(!hashing)entries.sort(Comparator.comparingInt(row->row.get("key").getAsInt()));
            node.add("entries",JSON.toJsonTree(entries));
        }else throw new IllegalArgumentException("Unsupported typed capture value: "+value.getClass().getName());
        return node;
    }
    private int visit(Value value,int depth) throws Exception {
        if(depth+heights.get(fingerprints.get(value))>maxDepth)throw new IllegalArgumentException("Typed capture exceeds depth bound");
        Integer old=identities.get(value);if(old!=null)return old;
        if(depth>maxDepth||identities.size()>=maxNodes)throw new IllegalArgumentException("Typed capture exceeds unique-node/depth bound");
        JsonObject node=describe(value,depth,false);String shape=JSON.toJson(node);Integer id=structures.get(shape);
        if(id==null){
            id=pool.size();if(id>=maxNodes)throw new IllegalArgumentException("Typed pool exceeds node bound");
            if(node.has("items"))edges+=node.getAsJsonArray("items").size();
            if(node.has("fields"))edges+=node.getAsJsonObject("fields").size();
            if(node.has("entries"))edges+=2L*node.getAsJsonArray("entries").size();
            if(edges>2000000)throw new IllegalArgumentException("Typed pool exceeds edge bound");
            structures.put(shape,id);pool.add(node);
        }identities.put(value,id);return id;
    }
    private Value[] reconstruct(){
        Value[] values=new Value[pool.size()];
        for(int id=0;id<pool.size();id++){
            JsonObject node=pool.get(id);String kind=node.get("kind").getAsString();Value value;
            if(kind.equals("integer"))value=IntValue.gen(Integer.parseInt(node.get("value").getAsString()));
            else if(kind.equals("boolean"))value=node.get("value").getAsBoolean()?BoolValue.ValTrue:BoolValue.ValFalse;
            else if(kind.equals("string"))value=new StringValue(node.get("value").getAsString());
            else if(kind.equals("model_value"))value=atoms.get(node.get("value").getAsString());
            else if(kind.equals("tuple")||kind.equals("set")){
                JsonArray items=node.getAsJsonArray("items");Value[] children=new Value[items.size()];
                for(int i=0;i<children.length;i++)children[i]=values[items.get(i).getAsInt()];
                value=kind.equals("tuple")?new TupleValue(children):new SetEnumValue(children,false);
            }else if(kind.equals("record")){
                JsonObject fields=node.getAsJsonObject("fields");UniqueString[] names=new UniqueString[fields.size()];Value[] children=new Value[fields.size()];int i=0;
                for(var entry:fields.entrySet()){names[i]=UniqueString.uniqueStringOf(entry.getKey());children[i]=values[entry.getValue().getAsInt()];i++;}
                value=new RecordValue(names,children,false);
            }else{
                JsonArray entries=node.getAsJsonArray("entries");Value[] domain=new Value[entries.size()],children=new Value[entries.size()];
                for(int i=0;i<children.length;i++){JsonObject row=entries.get(i).getAsJsonObject();domain[i]=values[row.get("key").getAsInt()];children[i]=values[row.get("value").getAsInt()];}
                value=new FcnRcdValue(domain,children,false);
            }if(value==null)throw new IllegalArgumentException("Absent reconstructed typed value");values[id]=value;
        }return values;
    }
    @tlc2.overrides.TLAPlusOperator(identifier="TraceEmit",module="FormalWikiTrace",warn=true)
    public static Value TraceEmit(Value value) throws Exception {
        String run=System.getProperty("formalwiki.run"),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();encoder.fingerprint(value,0);int root=encoder.visit(value,0);
        if(!value.equals(encoder.reconstruct()[root]))throw new IllegalArgumentException("Full typed capture reconstruction differs from TLC Value");
        JsonObject capture=new JsonObject();capture.addProperty("schema_version",2);capture.addProperty("kind","typed-tlc-capture");
        capture.addProperty("run_id",run);capture.addProperty("input_manifest_sha256",input);capture.addProperty("root",root);capture.add("nodes",JSON.toJsonTree(encoder.pool));
        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);}Files.createLink(target,temporary);}
        finally{Files.deleteIfExists(temporary);}return BoolValue.ValTrue;
    }
}