package org.bdware.sc.analysis.dynamic;

import org.bdware.analysis.BasicBlock;
import org.bdware.analysis.BreadthFirstSearch;
import org.bdware.analysis.taint.*;
import org.bdware.sc.bean.Contract;
import org.bdware.sc.compiler.YJSCompiler;
import org.bdware.sc.engine.DesktopEngine;
import org.bdware.sc.node.ContractNode;
import org.bdware.sc.node.ContractZipBundle;
import org.bdware.sc.node.FunctionNode;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.tree.ClassNode;
import org.objectweb.asm.tree.MethodNode;

import java.io.ByteArrayInputStream;
import java.io.FileInputStream;
import java.util.*;
import java.util.zip.ZipFile;

public class FSAnalysis extends BreadthFirstSearch<TaintResult, TaintBB> {
    TaintCFG cfg;
    public static boolean isDebug = false;

    public FSAnalysis(TaintCFG cfg) {
        this.cfg = cfg;
        List<TaintBB> toAnalysis = new ArrayList<>();
        // TODO add inputBlock!
        TaintBB b = (TaintBB) cfg.getBasicBlockAt(0);

        b.preResult = new TaintResult();
        // local0=scriptfuncion, is not tainted;
        // local1=this, is not tainted;
        // local2=this, is not tainted;

        b.preResult.frame.setLocal(0, HeapObject.getRootObject());
        b.preResult.frame.setLocal(1, new TaintValue(1, 0));
        b.preResult.frame.setLocal(2, new TaintValue(1, 1));

        b.preResult.ret = new TaintValue(1);
        TaintResult.printer.setLabelOrder(cfg.getLabelOrder());
        toAnalysis.add(b);
        b.setInList(true);
        setToAnalysis(toAnalysis);
        if (isDebug) {
            System.out.println("===Method:" + cfg.getMethodNode().name + cfg.getMethodNode().desc);
            System.out.println(
                    "===Local:"
                            + cfg.getMethodNode().maxLocals
                            + " "
                            + cfg.getMethodNode().maxStack);
        }
    }

    @Override
    public TaintResult execute(TaintBB t) {
        return t.forwardAnalysis();
    }

    @Override
    public Collection<TaintBB> getSuc(TaintBB t) {
        Set<BasicBlock> subBlock = cfg.getSucBlocks(t);
        Set<TaintBB> ret = new HashSet<>();
        for (BasicBlock bb : subBlock) {
            TaintBB ntbb = (TaintBB) bb;
            ntbb.preResult.mergeResult(t.sucResult);
            ret.add(ntbb);
        }

        return ret;
    }

    public static String staticVerify(Contract c) {
        try {
            String script = c.getScriptStr();
            ContractNode cn = null;
            YJSCompiler compiler = new YJSCompiler();
            if (script.startsWith("/")) {
                String zipPath = script;
                ZipFile zf = new ZipFile(zipPath);
                ContractZipBundle czb = compiler.compile(zf);
                cn = czb.mergeContractNode();
            } else {
                cn =
                        compiler.compile(
                                new ByteArrayInputStream(script.getBytes()), "contract_main.yjs");
            }
            DesktopEngine engine = new DesktopEngine(); // engine.loadJar(zf);
            engine.loadContract(c, cn, 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);
                }
            }
            Map<String, String> result = new HashMap<>();
            for (FunctionNode fn : cn.getFunctions()) {
                MethodNode mn = methods.get(fn.functionName);
                if (mn != null && mn.name.equals("statAge")) {
                    System.out.println("[ContractManager] verify:" + fn.functionName);
                    TaintResult.nLocals = mn.maxLocals;
                    TaintResult.nStack = mn.maxStack;
                    TaintCFG cfg = new TaintCFG(mn);
                    TaintResult.printer.setLabelOrder(cfg.getLabelOrder());
                    FSAnalysis analysis = new FSAnalysis(cfg);
                    analysis.analysis();
                    TaintBB bb = cfg.getLastBlock();
                    if (bb != null) result.put(fn.functionName, bb.getResult());
                    cfg.printSelf();
                }
            }
            return "success";
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }
}
