package com.android.commands.monkey.ape.jev;

import java.util.*;
import org.json.*;

/** Projection of accepted locator evidence into the existing shared-state / questions contract. */
public final class JevRegionEvidence {
    private JevRegionEvidence() { }
    /** Copy JSON containers without serializing and reparsing immutable strings. */
    private static Object copy(Object value)throws JSONException {
        if(value instanceof JSONObject)return copyObject((JSONObject)value);
        if(value instanceof JSONArray){JSONArray out=new JSONArray();JSONArray in=(JSONArray)value;
            for(int i=0;i<in.length();i++)out.put(copy(in.get(i)));return out;}
        return value;
    }
    private static JSONObject copyObject(JSONObject in,String... omitted)throws JSONException {
        Set<String> skip=omitted.length==0?Collections.<String>emptySet():new HashSet<String>(Arrays.asList(omitted));JSONObject out=new JSONObject();
        for(Iterator<?> it=in.keys();it.hasNext();){String k=(String)it.next();if(!skip.contains(k))out.put(k,copy(in.get(k)));}
        return out;
    }
    /** Wire-only factoring; the review evidence and its digest stay expanded. */
    static JSONObject wireNames(JSONObject state)throws JSONException {
        Map<String,Integer> counts=new TreeMap<String,Integer>();countNames(state,counts);
        Map<String,Integer> ids=new HashMap<String,Integer>();JSONArray names=new JSONArray();
        for(Map.Entry<String,Integer> entry:counts.entrySet())if(entry.getValue()>1){ids.put(entry.getKey(),names.length());names.put(entry.getKey());}
        if(ids.isEmpty())return state;
        return ((JSONObject)nameValues(state,ids,null)).put("full_names",names);
    }
    private static void countNames(Object value,Map<String,Integer> counts)throws JSONException {
        if(value instanceof JSONObject){JSONObject object=(JSONObject)value;
            for(Iterator<?> it=object.keys();it.hasNext();){String key=(String)it.next();Object item=object.get(key);
                if("full_name".equals(key) && item instanceof String){String name=(String)item;counts.put(name,counts.containsKey(name)?counts.get(name)+1:1);}
                else countNames(item,counts);}}
        else if(value instanceof JSONArray){JSONArray a=(JSONArray)value;for(int i=0;i<a.length();i++)countNames(a.get(i),counts);}
    }
    static JSONObject expandedNames(JSONObject state)throws JSONException {
        if(!state.has("full_names"))return state;JSONArray names=state.getJSONArray("full_names");
        for(int i=0;i<names.length();i++)if(!(names.get(i) instanceof String))throw new JSONException("invalid full_name dictionary");
        JSONObject root=new JSONObject();for(Iterator<?> it=state.keys();it.hasNext();){String k=(String)it.next();if(!"full_names".equals(k))root.put(k,state.get(k));}
        return (JSONObject)nameValues(root,null,names);
    }
    private static Object nameValues(Object value,Map<String,Integer> ids,JSONArray names)throws JSONException {
        if(value instanceof JSONObject){JSONObject in=(JSONObject)value,out=new JSONObject();
            for(Iterator<?> it=in.keys();it.hasNext();){String key=(String)it.next();Object item=in.get(key);
                if(ids!=null && "full_name".equals(key) && ids.containsKey(item)){out.put("full_name_ref",ids.get(item));}
                else if(names!=null && "full_name_ref".equals(key)){
                    if(in.has("full_name") || !(item instanceof Integer || item instanceof Long))throw new JSONException("invalid full_name reference");
                    long index=((Number)item).longValue();if(index<0 || index>=names.length())throw new JSONException("invalid full_name reference");
                    out.put("full_name",names.get((int)index));
                }else out.put(key,nameValues(item,ids,names));
            }return out;
        }
        if(value instanceof JSONArray){JSONArray in=(JSONArray)value,out=new JSONArray();for(int i=0;i<in.length();i++)out.put(nameValues(in.get(i),ids,names));return out;}
        return value;
    }
    /**
     * The complete locator receipt is already saved separately. The review FIFO
     * only needs the union of facts reachable by its nominated questions. Keep
     * those frozen values, not the entire UTG and every concrete's evidence.
     * Projecting any subset with state() must give the same request evidence.
     */
    static JSONObject pendingSnapshot(JSONObject snapshot,Set<Integer> requested)throws JSONException {
        JSONObject out=new JSONObject();
        for(String key:Arrays.asList("snapshot_monotonic_ns","run_start_monotonic_ns","elapsed_ms",
                "abstraction_version","topology_version","concrete_version","schedule"))out.put(key,snapshot.get(key));
        Set<Integer> internal=new HashSet<Integer>(),needed=new HashSet<Integer>();JSONArray regions=new JSONArray();
        JSONArray all=snapshot.getJSONArray("regions");
        for(int i=0;i<all.length();i++){
            JSONObject r=all.getJSONObject(i);if(!r.getBoolean("nominated"))continue;JSONArray entries=new JSONArray();
            for(int j=0;j<r.getJSONArray("entries").length();j++){
                JSONObject e=r.getJSONArray("entries").getJSONObject(j);
                if(!e.getBoolean("eligible_identity") || !requested.contains(e.getInt("concrete_id")))continue;
                entries.put(e);needed.add(e.getInt("source"));needed.addAll(ints(e.getJSONArray("known_targets")));
            }
            if(entries.length()==0)continue;
            JSONObject region=new JSONObject();for(Iterator<?> it=r.keys();it.hasNext();){String key=(String)it.next();if(!"entries".equals(key))region.put(key,r.get(key));}
            regions.put(region.put("entries",entries));internal.addAll(ints(r.getJSONArray("members")));
        }
        needed.addAll(internal);JSONArray transitions=new JSONArray(),history=snapshot.getJSONArray("transition_observations");
        for(int i=0;i<history.length();i++){
            JSONObject t=history.getJSONObject(i);
            if(!internal.contains(t.getInt("source")) && (t.isNull("concrete_id") || !requested.contains(t.getInt("concrete_id"))))continue;
            transitions.put(t);needed.add(t.getInt("source"));if(!t.isNull("target"))needed.add(t.getInt("target"));
        }
        JSONArray states=new JSONArray(),allStates=snapshot.getJSONArray("states");
        for(int i=0;i<allStates.length();i++)if(needed.contains(allStates.getJSONObject(i).getInt("index")))states.put(allStates.getJSONObject(i));
        return out.put("regions",regions).put("states",states).put("transition_observations",transitions);
    }
    public static JSONObject state(JSONObject snapshot,Map<Integer,JSONObject> details,Collection<Integer> selected)throws JSONException {
        Map<Integer,JSONObject> candidates=new TreeMap<Integer,JSONObject>();
        Map<Integer,JSONObject> regions=new TreeMap<Integer,JSONObject>();Set<Integer> needed=new TreeSet<Integer>();
        JSONArray all=snapshot.getJSONArray("regions"),history=snapshot.getJSONArray("transition_observations");
        Set<Integer> requested=new HashSet<Integer>(selected);
        for(int i=0;i<all.length();i++){
            JSONObject region=all.getJSONObject(i);if(!region.getBoolean("nominated"))continue;
            JSONArray entries=region.getJSONArray("entries");boolean used=false;
            for(int j=0;j<entries.length();j++){
                JSONObject entry=entries.getJSONObject(j);int id=entry.getInt("concrete_id");
                if(!requested.contains(id) || !entry.getBoolean("eligible_identity"))continue;
                JSONObject candidate=candidates.get(id);
                if(candidate==null){
                    JSONObject detail=details.get(id);if(detail==null)throw new JSONException("missing frozen concrete evidence: "+id);
                    candidate=copyObject(detail);candidate.put("candidate_id","c"+(candidates.size()+1));
                    candidate.put("known_targets",entry.getJSONArray("known_targets"));candidate.put("associated_regions",new JSONArray());
                    candidate.put("identity",JevRegistry.fields("concrete_id",id,"source",entry.getInt("source"),
                        "state_id",entry.getJSONObject("concrete_evidence").getString("state_id"),"model_action_id",entry.getJSONObject("concrete_evidence").getString("model_action_id"),
                        "source_gui_id",detail.opt("source_gui_id")));
                    candidate.remove("source_gui_id");candidate.put("history_indices",new JSONArray());
                    if(!candidate.has("evidence_gaps"))candidate.put("evidence_gaps",new JSONArray());
                    candidates.put(id,candidate);needed.add(entry.getInt("source"));
                    for(int k=0;k<entry.getJSONArray("known_targets").length();k++)needed.add(entry.getJSONArray("known_targets").getInt(k));
                }
                candidate.getJSONArray("associated_regions").put(JevRegistry.fields("header",region.getInt("header"),
                    "inside_targets",entry.getJSONArray("inside_targets"),"outside_targets",entry.getJSONArray("outside_targets"),"mixed_targets",entry.getBoolean("mixed_targets")));
                used=true;
            }
            if(used){
                JSONObject shared=copyObject(region,"entries");
                if(shared.getLong("state_progress_ns")==Long.MIN_VALUE)shared.put("state_progress_ns",JSONObject.NULL);
                regions.put(region.getInt("header"),shared);
                for(int j=0;j<region.getJSONArray("members").length();j++)needed.add(region.getJSONArray("members").getInt(j));
            }
        }
        if(candidates.size()!=requested.size())throw new JSONException("requested concrete missing from accepted nominations");
        Set<Integer> internal=new HashSet<Integer>();for(JSONObject region:regions.values())for(int i=0;i<region.getJSONArray("members").length();i++)internal.add(region.getJSONArray("members").getInt(i));
        JSONArray transitions=new JSONArray();
        for(int i=0;i<history.length();i++){
            JSONObject t=history.getJSONObject(i);JSONObject candidate=t.isNull("concrete_id")?null:candidates.get(t.getInt("concrete_id"));
            if(!internal.contains(t.getInt("source")) && candidate==null)continue;
            int index=transitions.length();transitions.put(copyObject(t));
            needed.add(t.getInt("source"));if(!t.isNull("target"))needed.add(t.getInt("target"));
            if(candidate!=null)candidate.getJSONArray("history_indices").put(index);
        }
        JSONArray states=new JSONArray(),sourceStates=snapshot.getJSONArray("states");
        for(int i=0;i<sourceStates.length();i++)if(needed.contains(sourceStates.getJSONObject(i).getInt("index"))){
            JSONObject s=copyObject(sourceStates.getJSONObject(i));
            // Native saturation is not a review prerequisite or a proxy for remaining value.
            s.remove("known");s.remove("saturated");s.remove("settled");states.put(s);
        }
        JSONArray cs=new JSONArray();for(JSONObject candidate:candidates.values())cs.put(candidate);
        // Assign packet-local IDs in exactly the order serialized, independent of traversal order.
        for(int i=0;i<cs.length();i++)cs.getJSONObject(i).put("candidate_id","c"+(i+1));
        return JevRegistry.fields("evidence_schema","ape_region_v1","snapshot_monotonic_ns",snapshot.getLong("snapshot_monotonic_ns"),
            "run_start_monotonic_ns",snapshot.getLong("run_start_monotonic_ns"),"elapsed_ms",snapshot.getLong("elapsed_ms"),
            "abstraction_version",snapshot.getLong("abstraction_version"),"topology_version",snapshot.getLong("topology_version"),
            "concrete_version",snapshot.getLong("concrete_version"),"regions",new JSONArray(regions.values()),"states",states,
            "transition_observations",transitions,"candidates",cs);
    }
    static void validate(JSONObject state)throws JSONException {
        JSONArray candidates=state.getJSONArray("candidates"),rs=state.getJSONArray("regions"),ss=state.getJSONArray("states"),ts=state.getJSONArray("transition_observations");
        Set<Integer> headers=new HashSet<Integer>(),states=new HashSet<Integer>(),concretes=new HashSet<Integer>();
        for(int i=0;i<ss.length();i++)if(!states.add(ss.getJSONObject(i).getInt("index")))throw new JSONException("duplicate State evidence");
        for(int i=0;i<rs.length();i++){
            JSONObject r=rs.getJSONObject(i);if(!headers.add(r.getInt("header")))throw new JSONException("duplicate Region evidence");
            JSONArray members=r.getJSONArray("members");Set<Integer> ms=ints(members);if(!ms.contains(r.getInt("header")) || !states.containsAll(ms))throw new JSONException("missing Region membership evidence");
        }
        for(int i=0;i<candidates.length();i++){
            JSONObject c=candidates.getJSONObject(i);int id=c.getJSONObject("identity").getInt("concrete_id");
            if(!concretes.add(id))throw new JSONException("duplicate concrete question");
            Set<Integer> targets=ints(c.getJSONArray("known_targets"));if(targets.isEmpty() || !states.containsAll(targets))throw new JSONException("missing target evidence");
            JSONArray regions=c.getJSONArray("associated_regions");if(regions.length()==0)throw new JSONException("missing candidate Region");
            for(int j=0;j<regions.length();j++){
                JSONObject r=regions.getJSONObject(j);Set<Integer> in=ints(r.getJSONArray("inside_targets")),out=ints(r.getJSONArray("outside_targets"));
                if(!headers.contains(r.getInt("header")) || !Collections.disjoint(in,out) || !in.contains(r.getInt("header")))throw new JSONException("invalid entry relation");
                in.addAll(out);if(!in.equals(targets))throw new JSONException("incomplete entry targets");
            }
            JSONArray indices=c.getJSONArray("history_indices");for(int j=0;j<indices.length();j++){
                int n=indices.getInt(j);if(n<0 || n>=ts.length() || ts.getJSONObject(n).getInt("concrete_id")!=id)throw new JSONException("invalid history reference");
            }
        }
    }
    private static Set<Integer> ints(JSONArray a)throws JSONException {Set<Integer> s=new HashSet<Integer>();for(int i=0;i<a.length();i++)s.add(a.getInt(i));return s;}
}