package org.bdware.analysis.gas;

import java.util.ArrayList;

import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.bdware.analysis.BasicBlock;
import org.bdware.analysis.CFGraph;
import org.bdware.analysis.OpInfo;
import org.bdware.analysis.gas.DFS;
import org.bdware.analysis.gas.FeeSchedule;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.ClassNode;
import org.objectweb.asm.tree.InvokeDynamicInsnNode;
import org.objectweb.asm.tree.JumpInsnNode;
import org.objectweb.asm.tree.LabelNode;

public class PPCount extends DFS {
    CFGraph cfg;
    String functionGlobel;
    private String bdName;
    private int globalInvoke = 0;
    static List<String> functionList = new ArrayList<>();
    private int globalBlockID;

    public static Map<String, HashMap<String, Integer>> varInsnCount = new HashMap<>();
    public static HashMap<String, Set<Map<Integer, HashMap<String, Integer>>>> branchCount =
            new HashMap<>();;
    public static Map<Integer, HashMap<String, Integer>> callFunction = new HashMap<>();
    public static Map<Integer, HashMap<String, Integer>> BlockInsn = new HashMap<>();
    public static Map<Integer, Set<Map<Integer, HashMap<String, Integer>>>> ppMap = new HashMap<>();
    public HashSet<Integer> countBlock = new HashSet<>();
    boolean[] isBranchBlock;

    public PPCount(CFGraph cfg, int flag) {
        this.cfg = cfg;
        functionGlobel = cfg.getMethodNode().name;
        varInsnCount.put(functionGlobel, new HashMap<>());
        globalInvoke = flag;
        ppMap.put(globalInvoke, new HashSet<>());
        callFunction.put(flag, new HashMap<>());
        List<BasicBlock> toAnalysis = new ArrayList<>();
        BasicBlock b = cfg.getBasicBlockAt(0);
        toAnalysis.add(b);
        b.setInList(true);
//        System.out.println("=========:"+cfg.getMethodNode().name);
        functionList.add(cfg.getMethodNode().name);
        isBranchBlock = new boolean[cfg.getBasicBlockSize()];
        // setToAnalysis(toAnalysis);
    }

    // @Override
    // public Collection<BasicBlock> getSuc(BasicBlock t) {
    // System.out.println("++++++++++" + t.blockID);
    // typeCount(t);
    // Set<BasicBlock> subBlock = cfg.getSucBlocks(t);
    // DFS dgDFS = new DFS(cfg, t);
    //
    // return subBlock;
    // }

    @Override
    public Set<BasicBlock> getSuc(BasicBlock t) {
        // System.out.println("++++++++++" + t.blockID);
        Set<BasicBlock> subBlock = cfg.getSucBlocks(t);
        typeCount(t);
        return subBlock;
    }

    private void typeCount(BasicBlock t) { // 一个基本块里面的所有指令
        int count = 0;
        countBlock.add(t.blockID);
        globalBlockID = t.blockID;
        BlockInsn.put(t.blockID, new HashMap<>());
        List<AbstractInsnNode> insnList = t.getInsn();
        OpInfo info = null;
        if (t.list.size() > 0) {
            for (AbstractInsnNode ab : insnList) {
                if (ab != null) {
                    if (ab.getOpcode() >= 0) {
                        count++;
                    }
                }
            }
        }
        t.insnCount = count;
        if (t.list.size() > 0) {
            for (int i = 0; i < insnList.size(); i++) {
                AbstractInsnNode insn = insnList.get(i);
                if (insn != null) {
                    if (insn.getOpcode() >= 0) {
                        info = OpInfo.ops[insn.getOpcode()];
                    }
                    if (info == null) {
                    } else if (info.canThrow()) {
                        callCount(insn, t);
                    } else if (info.canBranch()) {
                        normalCount(insn, t);
                        jumpCount(insn, t);
                    } else if (info.canContinue() || info.canSwitch() || info.canReturn()) {
                        normalCount(insn, t);
                    }
                }
            }
        }
        addBlock(t);
        removeBranch();
    }

    int jumpBlockid = -1;
    int conBlockid = -1;
    private boolean isBranch = false;
    private List<Integer> globalBranch = new ArrayList<>();

    private void jumpCount(AbstractInsnNode insn, BasicBlock t) {
        switch (insn.getOpcode()) {
            case Opcodes.IFNE: // succeeds if and only if value ≠ 0
            case Opcodes.IFEQ: // succeeds if and only if value = 0
            case Opcodes.IFLT: // succeeds if and only if value < 0
            case Opcodes.IFLE:

            case Opcodes.IFGE: // succeeds if and only if value ≥ 0
                isBranch = true;
                globalBranch.add(globalInvoke);
                Map<Integer, HashMap<String, Integer>> map = new HashMap<>();
                Map<Integer, HashMap<String, Integer>> map2 = new HashMap<>();
                Set<Map<Integer, HashMap<String, Integer>>> ret = new HashSet<>();
                Set<Map<Integer, HashMap<String, Integer>>> ret2 = new HashSet<>();
                map.put(t.blockID, BlockInsn.get(t.blockID));
                map2.put(t.blockID - 1, BlockInsn.get(t.blockID - 1));
                ret.add(map);
                ret.add(map2);
                ret2.add(map);
                ret2.add(map2);
                if (insn instanceof JumpInsnNode) {
                    LabelNode jump = ((JumpInsnNode) insn).label;
                    jumpBlockid = cfg.getBasicBlockByLabel(jump.getLabel()).blockID;

                    branchCount.put(globalInvoke + "false" + jumpBlockid, ret);
                }
                conBlockid = t.blockID + 1;

                branchCount.put(globalInvoke + "true" + conBlockid, ret2);
                // System.out.println(branchCount);
                break;
            case Opcodes.IFGT: // succeeds if and only if value > 0
                if (insn instanceof JumpInsnNode) {
                    LabelNode jump = ((JumpInsnNode) insn).label;
                    jumpBlockid = cfg.getBasicBlockByLabel(jump.getLabel()).blockID;
                    branchCount.put(globalInvoke + "false" + jumpBlockid, new HashSet<>());
                }
                conBlockid = t.blockID + 1;
                branchCount.put(globalInvoke + "true" + conBlockid, new HashSet<>());
                // System.out.println(branchCount);
                break;
            case Opcodes.IF_ACMPEQ: // succeeds if and only if value1 = value2
            case Opcodes.IF_ACMPNE: // succeeds if and only if value1 ≠ value2
            case Opcodes.IF_ICMPEQ: // succeeds if and only if value1 = value2
            case Opcodes.IF_ICMPGE: // succeeds if and only if value1 ≠ value2
            case Opcodes.IF_ICMPGT: // succeeds if and only if value1 > value2
            case Opcodes.IF_ICMPLE: // succeeds if and only if value1 ≤ value2
            case Opcodes.IF_ICMPLT: // succeeds if and only if value1 < value2
            case Opcodes.IFNONNULL:
            case Opcodes.TABLESWITCH:
            case Opcodes.LOOKUPSWITCH:
            case Opcodes.GOTO:
                lastBlockId = t.blockID + 1;
            default:
                break;
        }
    }

    int lastBlockId = 0;
    int globalBlock = 0;

    private void normalCount(AbstractInsnNode insn, BasicBlock t) {
        if (t.blockID != lastBlockId) {
            lastBlockId = t.blockID;
            globalBlock = 0;
        }
        OpInfo info = OpInfo.ops[insn.getOpcode()];
        bdName = FeeSchedule.BDInsn.name();
        if (BlockInsn.get(globalBlockID).containsKey(bdName)) {
            BlockInsn.get(globalBlockID).put(bdName, BlockInsn.get(globalBlockID).get(bdName) + 1);
        } else {
            BlockInsn.get(globalBlockID).put(bdName, 1);
        }
        globalBlock++;

        if (globalBranch.size() != 0) {
            int gb = globalBranch.get(globalBranch.size() - 1);
            if (gb == globalInvoke) {
                if (globalBlock == t.insnCount && !(insn instanceof JumpInsnNode)) {
                    chosenBranch(t);
                }
            }
        }
    }

    private void chosenBranch(BasicBlock t) {
        if (globalBranch.size() != 0) {
            int branch = globalBranch.get(globalBranch.size() - 1);
            Set<Map<Integer, HashMap<String, Integer>>> ret = new HashSet<>();
            if (isBranch && t.blockID < jumpBlockid) {
                isBranchBlock[t.blockID] = true;
                if (branchCount.get(branch + "true" + conBlockid) != null
                        && !branchCount.get(branch + "true" + conBlockid).isEmpty()) {
                    ret = branchCount.get(branch + "true" + conBlockid);
                    Map<Integer, HashMap<String, Integer>> map = new HashMap<>();
                    map.put(t.blockID, BlockInsn.get(t.blockID));
                    ret.add(map);
                    branchCount.put(branch + "true" + conBlockid, ret);

                } else if (branchCount.get(branch + "true" + conBlockid) != null) {
                    ret = branchCount.get(branch + "true" + conBlockid);
                    Map<Integer, HashMap<String, Integer>> map = new HashMap<>();
                    map.put(t.blockID, BlockInsn.get(t.blockID));
                    branchCount.get(branch + "true" + conBlockid).add(map);
                }
            } else if (isBranch && t.blockID >= jumpBlockid) {
                isBranchBlock[t.blockID] = true;
                if (branchCount.get(branch + "false" + jumpBlockid) != null
                        && !branchCount.get(branch + "false" + jumpBlockid).isEmpty()) {
                    Set<Map<Integer, HashMap<String, Integer>>> list =
                            branchCount.get(branch + "false" + jumpBlockid);
                    Map<Integer, HashMap<String, Integer>> map = new HashMap<>();
                    map.put(t.blockID, BlockInsn.get(t.blockID));
                    list.add(map);
                    branchCount.put(branch + "false" + jumpBlockid, list);

                } else if (branchCount.get(branch + "false" + jumpBlockid) != null) {
                    Map<Integer, HashMap<String, Integer>> map = new HashMap<>();
                    if (ret.isEmpty()) {
                        map.put(t.blockID, BlockInsn.get(t.blockID));
                        branchCount.get(branch + "false" + jumpBlockid).add(map);
                    }
                }
            }
        }
    }

    int perBlockId = 0;

    public void removeBranch() {
        if (!globalBranch.isEmpty()) {
            for (int i = 0; i < globalBranch.size(); i++) {
                if (ppMap.containsKey(globalBranch.get(i))) {
                    ppMap.remove(globalBranch.get(i));
                }
            }
        }
    }

    private void callCount(AbstractInsnNode insn, BasicBlock t) {
        boolean isNewCall = false;
        lastBlockId = t.blockID;
        if (insn instanceof InvokeDynamicInsnNode) {
            Object invoke = ((InvokeDynamicInsnNode) insn).bsmArgs[0];
            String functionName = ((InvokeDynamicInsnNode) insn).name;
            isNewCall = true;
            if ((int) invoke > 5) {
                dynCount(functionName);
                globalInvoke = (int) invoke;
                ppMap.put((int) invoke, new HashSet<>());
                perBlockId = t.blockID;
            } else {
                dynCount(functionName);
            }
        } else {
            normalCount(insn, t);
        }
        if (!isNewCall || t.blockID == jumpBlockid || t.blockID == conBlockid) {
            branchEnd = true;
            chosenBranch(t);
        }
    }

    private void addBlock(BasicBlock t) {
        if (t.list.size() > 0) {
            Set<Map<Integer, HashMap<String, Integer>>> ret = new HashSet<>();
            Map<Integer, HashMap<String, Integer>> map = new HashMap<>();
            map.put(t.blockID, BlockInsn.get(t.blockID));
            if (!isBranchBlock[t.blockID]) {
                if (!ppMap.containsKey(globalInvoke)) {
                    ppMap.put(globalInvoke, new HashSet<>());
                }
                ppMap.get(globalInvoke).add(map);
            }
        }
    }

    boolean branchEnd;

    private void dynCount(String functionName) {
        String function = functionName.split(":")[1];
        if (function.contains("getProp")) {
            bdName = FeeSchedule.BDgetMethod.name();
            if (BlockInsn.get(globalBlockID).containsKey(bdName)) {
                BlockInsn.get(globalBlockID)
                        .put(bdName, BlockInsn.get(globalBlockID).get(bdName) + 1);
            } else {
                BlockInsn.get(globalBlockID).put(bdName, 1);
            }
        } else if (function.contains("setProp")) {
            bdName = FeeSchedule.BDsetMethod.name();
            if (BlockInsn.get(globalBlockID).containsKey(bdName)) {
                BlockInsn.get(globalBlockID)
                        .put(bdName, BlockInsn.get(globalBlockID).get(bdName) + 1);
            } else {
                BlockInsn.get(globalBlockID).put(bdName, 1);
            }
        } else if (function.contains("new")) {
            bdName = FeeSchedule.BDnew.name();
            if (BlockInsn.get(globalBlockID).containsKey(bdName)) {
                BlockInsn.get(globalBlockID)
                        .put(bdName, BlockInsn.get(globalBlockID).get(bdName) + 1);
            } else {
                BlockInsn.get(globalBlockID).put(bdName, 1);
            }
        } else if (function.contains("call") && functionName.split(":")[2].contains("Util")) {
            bdName = FeeSchedule.BDcallUtil.name();
            if (BlockInsn.get(globalBlockID).containsKey(bdName)) {
                BlockInsn.get(globalBlockID)
                        .put(bdName, BlockInsn.get(globalBlockID).get(bdName) + 1);
            } else {
                BlockInsn.get(globalBlockID).put(bdName, 1);
            }
        } else if (function.contains("call") && isInnerfunction(functionName)) {
           // System.out.println("[call ]" + functionName);
            String methName = functionName.split(":")[2];
            bdName = FeeSchedule.BDcallFuntion.name();
            if (BlockInsn.get(globalBlockID).containsKey(bdName)) {
                BlockInsn.get(globalBlockID)
                        .put(
                                bdName + "," + methName,
                                BlockInsn.get(globalBlockID).get(bdName + "," + methName) + 1);
                System.out.println("   function   " + functionName);
            } else {
                BlockInsn.get(globalBlockID).put(bdName + "," + methName, 1);
            }

        } else if (function.contains("call")) {
            bdName = FeeSchedule.BDcall.name();
            if (BlockInsn.get(globalBlockID).containsKey(bdName)) {
                BlockInsn.get(globalBlockID)
                        .put(bdName, BlockInsn.get(globalBlockID).get(bdName) + 1);
            } else {
                BlockInsn.get(globalBlockID).put(bdName, 1);
            }
        } else if (functionName.contains("traceif")) {
            bdName = FeeSchedule.BDjump.name();
            if (BlockInsn.get(globalBlockID).containsKey(bdName)) {
                BlockInsn.get(globalBlockID)
                        .put(bdName, BlockInsn.get(globalBlockID).get(bdName) + 1);
            } else {
                BlockInsn.get(globalBlockID).put(bdName, 1);
            }
        }
    }

    private boolean isInnerfunction(String functionName) {
        String string = functionName.split(":")[2];
       // System.out.println("++++++++++++++" + functionName);
       // System.out.println("【function】" + functionList);
        if (functionList.contains(string)) {
            return true;
        } else {
            return false;
        }
    }

    //    public static void main(String[] args) throws Exception {
    //        String path = "/Users/hulingxuan/git/SmartContract/output/main.yjs";
    //        ContractNode contractNode = null;
    //        YJSCompiler compiler = new YJSCompiler();
    //        contractNode = compiler.compile(new FileInputStream(path), null);
    //        engine = new DesktopEngine();
    //        engine.loadContract(contractNode, false);
    //
    //        Map<String, byte[]> clzs = engine.dumpClass();
    //        Map<String, MethodNode> methods = new HashMap<>();
    //        for (byte[] clz : clzs.values()) {
    //            ClassNode classNode = new ClassNode();
    //            ClassReader cr = new ClassReader(clz);
    //            cr.accept(classNode, ClassReader.EXPAND_FRAMES);
    //            for (MethodNode mn : classNode.methods) {
    //                methods.put(mn.name, mn);
    //            }
    //        }
    //        int flag = 0;
    //        for (FunctionNode fn : contractNode.getFunctions()) {
    //            functionList.add(fn.functionName);
    //            MethodNode mn = methods.get(fn.functionName);
    //            if (mn != null) {
    //                CFGraph cfg = new CFGraph(mn) {
    //                    @Override
    //                    public BasicBlock getBasicBlock(int id) {
    //                        return new BasicBlock(id);
    //                    }
    //                };
    //                // cfg.getLabelOrder();
    //                PPCount countFee = new PPCount(cfg, flag);
    //
    //                BasicBlock bb = cfg.getBasicBlockAt(0);
    //                countFee.dfs(cfg, bb);
    //                // cfg.printSelf();
    //
    //                Evaluates feEvaluates = new Evaluates();
    //                feEvaluates.getGas(branchCount);
    //                feEvaluates.getInsnGas(ppMap);
    //                countFunction(fn.functionName, Evaluates.map);
    //                System.out.println("+++++++" + functionSumGas);
    //                System.out.println("========" + Evaluates.map);
    //                flag++;
    //            }
    //        }
    //        System.out.println(branchCount);
    //        System.out.println(BlockInsn);
    //        System.out.println(ppMap);
    //
    //        // System.out.println(feEvaluates.getCallFee());
    //    }

    public static HashMap<String, Long> functionSumGas = new HashMap<>();

    public static void countFunction(String functionName, HashMap<String, Long> map) {
        long sumTrue = 0;
        long sumFalse = 0;
        long sum = 0;

        for (Map.Entry<String, Long> entry : map.entrySet()) {
            if (entry.getKey().contains("true")) {
                sumTrue = entry.getValue();
            } else if (entry.getKey().contains("false")) {
                sumFalse = entry.getValue();
            } else {
                sum += entry.getValue();
            }
        }
        functionSumGas.put(functionName, sum);
        functionSumGas.put(functionName + "true", sumTrue);
        functionSumGas.put(functionName + "false", sumFalse);
    }
}
