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

import android.graphics.*;
import android.util.Base64;
import java.io.*;
import java.util.*;
import org.json.*;
import com.android.commands.monkey.ape.tree.*;
import com.android.commands.monkey.ape.utils.Config;

/** Optional immutable visual evidence from the same bitmap used by native APE. */
public final class JevImages {
    private JevImages() {}
    private static long sequence;
    private static volatile long storedBytes;
    private static final int MAX_IMAGE_BYTES=128*1024;
    public static final class Snapshot implements Serializable {
        private static final long serialVersionUID=1L;
        public final String path;
        public final int originalWidth,originalHeight;
        Snapshot(String p,int w,int h){path=p;originalWidth=w;originalHeight=h;}
    }
    private static JevOutput output;
    private static final java.util.concurrent.Semaphore imageSlots=new java.util.concurrent.Semaphore(2);
    static void setOutput(JevOutput writer){output=writer;}
    public static Snapshot capture(Bitmap bitmap) {
        if(output==null || bitmap==null || bitmap.isRecycled() || storedBytes>=512L*1024*1024 || !imageSlots.tryAcquire())return null;
        Bitmap copy=null;
        try {
            final int w=bitmap.getWidth(),h=bitmap.getHeight();
            copy=bitmap.copy(Bitmap.Config.ARGB_8888,false);final Bitmap owned=copy;
            final File dir=new File(Config.get("ape.outputRoot","/sdcard/ape"),"jev-images");
            final File file=new File(dir,"frame-"+(++sequence)+".jpg");
            boolean accepted=output.submit("image_capture_compress_io",new JevOutput.Write(){public void run()throws Exception{
                Bitmap small=null;
                try {
                    float scale=Math.min(1f,800f/Math.max(w,h));
                    small=Bitmap.createScaledBitmap(owned,Math.max(1,Math.round(w*scale)),Math.max(1,Math.round(h*scale)),true);
                    byte[] data=jpeg(small);if(data==null || storedBytes+data.length>512L*1024*1024)throw new IOException("Image budget exceeded");
                    if(!dir.isDirectory() && !dir.mkdirs())throw new IOException("Image directory unavailable");
                    FileOutputStream stream=new FileOutputStream(file);try{stream.write(data);}finally{stream.close();}storedBytes+=data.length;
                } finally {if(small!=null && small!=owned)small.recycle();owned.recycle();imageSlots.release();}
            }});
            if(!accepted){owned.recycle();imageSlots.release();return null;}
            return new Snapshot(file.getAbsolutePath(),w,h);
        } catch(Exception e){if(copy!=null)copy.recycle();imageSlots.release();output.fail(e);return null;}
    }
    private static byte[] jpeg(Bitmap bitmap) {
        for(int quality:new int[]{85,65,45}) {
            ByteArrayOutputStream out=new ByteArrayOutputStream();
            if(!bitmap.compress(Bitmap.CompressFormat.JPEG,quality,out))return null;
            if(out.size()<=MAX_IMAGE_BYTES)return out.toByteArray();
        }
        return null;
    }
    private static final class Panel {
        final Snapshot snapshot; final JSONArray marks=new JSONArray(),refs=new JSONArray();
        Panel(Snapshot s){snapshot=s;}
    }
    private static void add(LinkedHashMap<String,Panel> panels,Snapshot snapshot,String label,JSONArray bounds,boolean candidate) throws JSONException {
        if(snapshot==null)return;
        Panel panel=panels.get(snapshot.path);
        if(panel==null){panel=new Panel(snapshot);panels.put(snapshot.path,panel);}
        panel.refs.put(label);
        if(bounds!=null)panel.marks.put(new JSONObject().put("label",label).put("bounds",bounds).put("candidate",candidate));
    }
    private static GUITreeTransition fact(JSONObject member,Map<String,GUITreeTransition> facts) throws JSONException {
        JSONArray ids=member.getJSONArray("transition_ids");
        for(int i=0;i<ids.length();i++){GUITreeTransition t=facts.get(ids.getString(i));if(t!=null)return t;}
        return null;
    }
    public static java.util.concurrent.Callable<JevProtocol.Request> prepare(final JevProtocol.Request request,Snapshot current,JSONObject receipt,
            Map<String,GUITreeTransition> facts,final int maxBytes) throws JSONException {
        JSONObject state=new JSONObject(request.stateJson);
        LinkedHashMap<String,Panel> panels=new LinkedHashMap<String,Panel>();
        JSONArray missing=new JSONArray();
        if(current==null)throw new JSONException("CURRENT_IMAGE_UNAVAILABLE");
        // Keep current capture first even if all bounds are absent.
        add(panels,current,"current GUI",null,false);
        for(String section:new String[]{"candidates","current_siblings"}) {
            JSONArray members=state.getJSONArray(section);
            String field=section.equals("candidates")?"candidate_id":"sibling_id";
            for(int i=0;i<members.length();i++) {
                JSONObject m=members.getJSONObject(i);
                add(panels,current,m.getString(field),m.optJSONArray("bounds"),field.equals("candidate_id"));
            }
        }
        JSONArray histories=receipt.getJSONArray("executed_history_other_pages");
        for(int i=0;i<histories.length();i++) {
            JSONObject m=histories.getJSONObject(i);String id=m.getString("member_id");
            GUITreeTransition t=fact(m,facts);Snapshot s=t==null?null:t.getSource().getJevImage();
            if(s!=null && (panels.containsKey(s.path) || panels.size()<3)) {
                android.graphics.Rect b=t.getAction().getGUITreeNode().getBoundsInScreen();
                add(panels,s,id,new JSONArray().put(b.left).put(b.top).put(b.right).put(b.bottom),false);
            } else missing.put(id+" source image omitted or unavailable");
        }
        for(String section:new String[]{"executed_history_other_pages","current_siblings"}) {
            JSONArray members=receipt.getJSONArray(section);
            String field=section.equals("current_siblings")?"sibling_id":"member_id";
            for(int i=0;i<members.length();i++) {
                JSONObject m=members.getJSONObject(i);String id=m.getString(field);
                JSONArray ids=m.getJSONArray("transition_ids");
                // Only associate a result picture with a sampled outcome actually in the text.
                JSONArray views=state.getJSONArray(section);JSONObject view=null;
                for(int j=0;j<views.length();j++)if(id.equals(views.getJSONObject(j).getString(field)))view=views.getJSONObject(j);
                if(view==null)continue;
                JSONArray outcomes=view.getJSONArray("outcomes");
                for(int j=0;j<outcomes.length();j++) {
                    JSONObject outcome=outcomes.getJSONObject(j);GUITreeTransition match=null;
                    for(int k=0;k<ids.length();k++) {
                        GUITreeTransition t=facts.get(ids.getString(k));
                        if(t!=null && outcome.getString("target_state").equals(t.getTarget().getCurrentState().getGraphId())
                                && outcome.getJSONArray("result_labels").toString().equals(JevEvidence.resultLabels(t.getTarget().getRootNode()).toString())){match=t;break;}
                    }
                    Snapshot s=match==null?null:match.getTarget().getJevImage();
                    String label=id+" outcome["+j+"] target "+outcome.getString("target_state");
                    if(s!=null && (panels.containsKey(s.path) || panels.size()<4))add(panels,s,label,null,false);
                    else missing.put(label+" image omitted or unavailable");
                }
            }
        }
        JSONArray frozenPanels=new JSONArray();
        for(Panel panel:panels.values())frozenPanels.put(new JSONObject().put("path",panel.snapshot.path)
            .put("width",panel.snapshot.originalWidth).put("height",panel.snapshot.originalHeight).put("marks",panel.marks).put("refs",panel.refs));
        final Object frozen=JevOutput.freeze(frozenPanels),frozenMissing=JevOutput.freeze(missing);
        return new java.util.concurrent.Callable<JevProtocol.Request>() {public JevProtocol.Request call()throws Exception {
            LinkedHashMap<String,Panel> panels=new LinkedHashMap<String,Panel>();
            JSONArray list=(JSONArray)JevOutput.json(frozen);JSONArray missing=(JSONArray)JevOutput.json(frozenMissing);
            for(int i=0;i<list.length();i++) {
                JSONObject item=list.getJSONObject(i);Panel panel=new Panel(new Snapshot(item.getString("path"),item.getInt("width"),item.getInt("height")));
                JSONArray marks=item.getJSONArray("marks"),refs=item.getJSONArray("refs");
                for(int j=0;j<marks.length();j++)panel.marks.put(marks.get(j));for(int j=0;j<refs.length();j++)panel.refs.put(refs.get(j));
                panels.put(panel.snapshot.path,panel);
            }
            return render(request,panels,missing,maxBytes);
        }};
    }
    private static JevProtocol.Request render(JevProtocol.Request request,LinkedHashMap<String,Panel> panels,JSONArray missing,int maxBytes)throws JSONException {
        JSONArray images=new JSONArray(),descriptions=new JSONArray();
        for(Panel panel:panels.values()) {
            Bitmap original=BitmapFactory.decodeFile(panel.snapshot.path),marked=null;
            if(original==null){missing.put(panel.refs.toString()+" image file unavailable");continue;}
            try {
                marked=original.copy(Bitmap.Config.ARGB_8888,true);
                int width=marked.getWidth(),height=marked.getHeight();
                int[] pixels=new int[width*height];marked.getPixels(pixels,0,width,0,0,width,height);
                float sx=marked.getWidth()/(float)panel.snapshot.originalWidth,sy=marked.getHeight()/(float)panel.snapshot.originalHeight;
                // Draw member outlines before candidate outlines so candidates stay visible on overlap.
                for(boolean candidates:new boolean[]{false,true})for(int i=0;i<panel.marks.length();i++) {
                    JSONObject mark=panel.marks.getJSONObject(i);if(mark.getBoolean("candidate")!=candidates)continue;
                    JSONArray b=mark.getJSONArray("bounds");float left=(float)b.getDouble(0)*sx,top=(float)b.getDouble(1)*sy,right=(float)b.getDouble(2)*sx,bottom=(float)b.getDouble(3)*sy;
                    int color=candidates?0xffff8c00:0xff00ffff;
                    int l=Math.round(left),t=Math.round(top),r=Math.round(right),bott=Math.round(bottom),stroke=candidates?3:2;
                    rect(pixels,width,height,l,t,r,t+stroke,color);rect(pixels,width,height,l,bott-stroke,r,bott,color);
                    rect(pixels,width,height,l,t,l+stroke,bott,color);rect(pixels,width,height,r-stroke,t,r,bott,color);
                    label(pixels,width,height,mark.getString("label"),Math.max(0,l+2),Math.max(0,Math.min(height-18,t+(candidates?2:20))),color);
                }
                marked.setPixels(pixels,0,width,0,0,width,height);
                byte[] bytes=jpeg(marked);if(bytes==null){missing.put(panel.refs.toString()+" annotation exceeds image budget");continue;}
                images.put("data:image/jpeg;base64,"+Base64.encodeToString(bytes,Base64.NO_WRAP));
                descriptions.put(new JSONObject().put("image","<image>").put("references",panel.refs)
                    .put("sha256",sha(bytes)).put("width",marked.getWidth()).put("height",marked.getHeight())
                    .put("original_width",panel.snapshot.originalWidth).put("original_height",panel.snapshot.originalHeight));
            } finally{if(marked!=null)marked.recycle();original.recycle();}
        }
        if(images.length()==0 || !descriptions.getJSONObject(0).getJSONArray("references").getString(0).equals("current GUI"))
            throw new JSONException("CURRENT_IMAGE_UNAVAILABLE");
        JSONObject visual=new JSONObject().put("panels",descriptions).put("omitted",missing).put("legend",
            "Panels are contemporaneous APE captures. Orange c# marks candidates; cyan s#/h# marks comparison members. Overlapping c#/s# may denote the same current control. Outcome references identify confirmed result pages. Bounds in member records use original screen coordinates; overlays are scaled to the displayed image. Listed omitted images have text evidence only.");
        return JevProtocol.withImages(request,visual,images,maxBytes);
    }
    // app_process has no app font initialization. Tiny pixel IDs avoid all font/Canvas APIs.
    private static void rect(int[] pixels,int w,int h,int left,int top,int right,int bottom,int color) {
        for(int y=Math.max(0,top);y<Math.min(h,bottom);y++)
            for(int x=Math.max(0,left);x<Math.min(w,right);x++)pixels[y*w+x]=color;
    }
    private static void label(int[] pixels,int w,int h,String text,int left,int top,int color) {
        String alphabet="0123456789csh";
        String[] glyphs={"111101101101111","010110010010111","111001111100111","111001111001111",
            "101101111001001","111100111001111","111100111101111","111001001001001",
            "111101111101111","111101111001111","111100100100111","111100111001111","101101111101101"};
        rect(pixels,w,h,left,top,left+text.length()*12+2,top+17,0xff000000);
        for(int i=0;i<text.length();i++) {
            int index=alphabet.indexOf(text.charAt(i));if(index<0)continue;
            String glyph=glyphs[index];
            for(int y=0;y<5;y++)for(int x=0;x<3;x++)if(glyph.charAt(y*3+x)=='1')
                rect(pixels,w,h,left+i*12+x*3+1,top+y*3+1,left+i*12+x*3+4,top+y*3+4,color);
        }
    }
    private static String sha(byte[] bytes) {
        try {StringBuilder hex=new StringBuilder();for(byte b:java.security.MessageDigest.getInstance("SHA-256").digest(bytes))hex.append(String.format(java.util.Locale.ROOT,"%02x",b&255));return hex.toString();}
        catch(java.security.NoSuchAlgorithmException e){throw new AssertionError(e);}
    }
}