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

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

/** Transport partitioning only: stable candidate IDs and complete execution evidence. */
public final class JevProbePackets {
    private JevProbePackets() {}
    private static final String[] TABLES={"role_signatures","page_labels","outcomes","local_contexts","source_snapshots","result_label_sets"};
    private static final String[] NO_OUTCOME_TABLES={"role_signatures","page_labels","local_contexts"};

    public static final class Plan {
        public final JevProtocol.Request logical;
        public final int hardLimitBytes,packetLimitBytes;
        public final List<JevProtocol.Request> packets;
        private Plan(JevProtocol.Request logical,int hard,int target,List<JevProtocol.Request> packets) {
            this.logical=logical;hardLimitBytes=hard;packetLimitBytes=target;
            this.packets=Collections.unmodifiableList(new ArrayList<JevProtocol.Request>(packets));
        }
        public JSONObject receipt()throws JSONException {
            JSONArray rows=new JSONArray();
            for(int i=0;i<packets.size();i++) {
                JevProtocol.Request p=packets.get(i);
                rows.put(new JSONObject().put("packet_index",i+1).put("physical_request_id",p.requestId)
                    .put("candidate_ids",new JSONArray(p.candidateIds)).put("request_bytes",p.requestBytes).put("evidence_digest",p.evidenceDigest));
            }
            return new JSONObject().put("logical_request_id",logical.requestId).put("logical_request_bytes",logical.requestBytes)
                .put("logical_evidence_digest",logical.evidenceDigest).put("hard_limit_bytes",hardLimitBytes)
                .put("packet_limit_bytes",packetLimitBytes).put("packets",rows);
        }
    }

    /** Reserve 4 KiB at the normal 64 KiB cap; scale the margin for smaller test/config budgets. */
    public static int packetLimit(int hardLimit) {
        if(hardLimit<=0)throw new IllegalArgumentException("invalid packet budget");
        return hardLimit-Math.min(4096,hardLimit/16);
    }

    public static Plan plan(JevProtocol.Request logical,int hardLimit)throws JSONException {
        if(logical.semantic)return semanticPlan(logical,hardLimit);
        Slice slice=new Slice(logical);int effectiveHard=slice.noOutcome?Math.min(hardLimit,32768):hardLimit;
        int target=packetLimit(effectiveHard);
        List<JevProtocol.Request> packets=new ArrayList<JevProtocol.Request>();
        List<String> current=new ArrayList<String>();
        JSONArray reviews=slice.state.getJSONArray("reviews");
        for(int i=0;i<reviews.length();i++) {
            List<String> group=new ArrayList<String>();
            JSONArray candidates=reviews.getJSONObject(i).getJSONArray("candidates");
            for(int j=0;j<candidates.length();j++)group.add(candidates.getJSONObject(j).getString("candidate_id"));
            List<String> combined=new ArrayList<String>(current);combined.addAll(group);
            if(slice.bytes(combined)<=target) {current=combined;continue;}
            if(!current.isEmpty()) {packets.add(slice.finish(current,packets.size()+1,target));current.clear();}
            if(slice.bytes(group)<=target) {current=group;continue;}
            // A large ModelAction is split by candidates, never by history member or outcome.
            for(String candidate:group) {
                combined=new ArrayList<String>(current);combined.add(candidate);
                if(slice.bytes(combined)<=target) {current=combined;continue;}
                if(!current.isEmpty()) {packets.add(slice.finish(current,packets.size()+1,target));current.clear();}
                current.add(candidate);
                if(slice.bytes(current)>target) slice.finish(current,packets.size()+1,target); // Explicit BudgetExceeded; no truncation.
            }
        }
        if(!current.isEmpty())packets.add(slice.finish(current,packets.size()+1,target));
        if(packets.isEmpty())throw new JSONException("empty logical probe request");
        List<String> assigned=new ArrayList<String>();for(JevProtocol.Request p:packets)assigned.addAll(p.candidateIds);
        if(!assigned.equals(logical.candidateIds))throw new JSONException("packet candidate partition changed");
        return new Plan(logical,effectiveHard,target,packets);
    }

    private static Plan semanticPlan(JevProtocol.Request logical,int hardLimit)throws JSONException {
        int target=packetLimit(hardLimit);List<JevProtocol.Request> packets=new ArrayList<JevProtocol.Request>();
        List<String> current=new ArrayList<String>();
        for(String id:logical.candidateIds) {
            List<String> trial=new ArrayList<String>(current);trial.add(id);
            try { JevProtocol.semanticPacket(logical,trial,packets.size()+1,target); current=trial; }
            catch(JevProtocol.BudgetExceeded tooLarge) {
                if(current.isEmpty())throw tooLarge;
                packets.add(JevProtocol.semanticPacket(logical,current,packets.size()+1,target));
                current=new ArrayList<String>();current.add(id);
                JevProtocol.semanticPacket(logical,current,packets.size()+1,target);
            }
        }
        if(!current.isEmpty())packets.add(JevProtocol.semanticPacket(logical,current,packets.size()+1,target));
        if(packets.isEmpty())throw new JSONException("empty logical semantic request");
        return new Plan(logical,hardLimit,target,packets);
    }

    /** Exact standalone packet size, including its full rules, context, history and question. */
    public static JevProtocol.Request single(JevProtocol.Request logical,String candidate,int maxBytes)throws JSONException {
        return new Slice(logical).finish(Collections.singletonList(candidate),1,maxBytes);
    }

    private static final class Slice {
        final JevProtocol.Request logical;
        final JSONObject body,state,shared,contexts,questions;
        final boolean noOutcome;
        Slice(JevProtocol.Request logical)throws JSONException {
            this.logical=logical;body=new JSONObject(logical.bodyJson);state=body.getJSONObject("state");
            String schema=state.getString("evidence_schema");
            noOutcome="ape_probe_v4".equals(schema);
            if(!noOutcome && !"ape_probe_v3".equals(schema))throw new JSONException("invalid packet evidence schema");
            shared=state.getJSONObject("shared_evidence");contexts=state.getJSONObject("history_contexts");questions=body.getJSONObject("questions");
        }
        int bytes(List<String> ids)throws JSONException{return subset(ids).toString().getBytes(JevProtocol.UTF8).length;}
        JevProtocol.Request finish(List<String> ids,int index,int limit)throws JSONException {
            return JevProtocol.wire(logical.requestId+":p"+index,subset(ids).toString(),limit,body.getString("model"));
        }
        JSONObject subset(List<String> ids)throws JSONException {
            Set<String> selected=new LinkedHashSet<String>(ids);
            if(selected.isEmpty() || selected.size()!=ids.size() || !logical.candidateIds.containsAll(ids))throw new JSONException("invalid packet candidate IDs");
            JSONObject out=new JSONObject();for(String k:JevSharedEvidence.keys(state))out.put(k,state.get(k));
            JSONArray reviews=new JSONArray();JSONArray original=state.getJSONArray("reviews");
            for(int i=0;i<original.length();i++) {
                JSONObject review=original.getJSONObject(i);JSONArray candidates=review.getJSONArray("candidates"),kept=new JSONArray();
                for(int j=0;j<candidates.length();j++)if(selected.contains(candidates.getJSONObject(j).getString("candidate_id")))kept.put(candidates.get(j));
                if(kept.length()==0)continue;
                JSONObject copy=new JSONObject();for(String k:JevSharedEvidence.keys(review))copy.put(k,review.get(k));
                copy.put("candidates",kept);reviews.put(copy);
            }
            out.put("reviews",reviews);
            Closure closure=new Closure(shared,contexts,noOutcome);closure.visit(state.get("source_context"));closure.visit(reviews);closure.resolve();
            JSONObject retained=new JSONObject();for(String name:closure.names)retained.put(name,closure.table(name));
            out.put("shared_evidence",retained).put("history_contexts",closure.table("history_contexts"));
            JSONObject qs=new JSONObject();for(String id:logical.candidateIds)if(selected.contains(id))qs.put(id,questions.get(id));
            return new JSONObject().put("model",body.get("model")).put("state",out).put("questions",qs);
        }
    }

    private static final class Closure {
        final Map<String,JSONObject> tables=new LinkedHashMap<String,JSONObject>();
        final String[] names;
        final Map<String,Set<String>> used=new LinkedHashMap<String,Set<String>>();
        final Deque<String[]> pending=new ArrayDeque<String[]>();
        Closure(JSONObject shared,JSONObject contexts,boolean noOutcome)throws JSONException {
            names=noOutcome?NO_OUTCOME_TABLES:TABLES;
            for(String name:names)tables.put(name,shared.getJSONObject(name));tables.put("history_contexts",contexts);
            for(String name:tables.keySet())used.put(name,new HashSet<String>());
        }
        void reference(String table,Object id)throws JSONException {
            if(!(id instanceof String) || !tables.get(table).has((String)id))throw new JSONException("missing packet evidence reference in "+table);
            if(used.get(table).add((String)id))pending.add(new String[]{table,(String)id});
        }
        void visit(Object value)throws JSONException {
            if(value instanceof JSONArray) {
                JSONArray a=(JSONArray)value;for(int i=0;i<a.length();i++)visit(a.get(i));
            } else if(value instanceof JSONObject) {
                JSONObject o=(JSONObject)value;
                for(String key:JevSharedEvidence.keys(o)) {
                    Object v=o.get(key);
                    if("role_signature".equals(key))reference("role_signatures",v);
                    else if("source_snapshot".equals(key))reference("source_snapshots",v);
                    else if("source_context_id".equals(key))reference("history_contexts",v);
                    else if("result_labels".equals(key))reference("result_label_sets",v);
                    else if("local_context".equals(key)) {if(v!=JSONObject.NULL)reference("local_contexts",v);}
                    else if("outcomes".equals(key) || "page_labels".equals(key)) {
                        JSONArray a=o.getJSONArray(key);for(int i=0;i<a.length();i++)reference(key,a.get(i));
                    } else visit(v);
                }
            }
        }
        void resolve()throws JSONException {
            while(!pending.isEmpty()) {String[] next=pending.remove();visit(tables.get(next[0]).get(next[1]));}
        }
        JSONObject table(String name)throws JSONException {
            JSONObject result=new JSONObject(),source=tables.get(name);
            for(String id:JevSharedEvidence.keys(source))if(used.get(name).contains(id))result.put(id,source.get(id));
            return result;
        }
    }
}