package com.ibm.wala.shrikeCT;

import static com.ibm.wala.shrikeBT.Constants.TYPE_double;
import static com.ibm.wala.shrikeBT.Constants.TYPE_float;
import static com.ibm.wala.shrikeBT.Constants.TYPE_int;
import static com.ibm.wala.shrikeBT.Constants.TYPE_long;
import static com.ibm.wala.shrikeBT.Constants.TYPE_null;

import com.ibm.wala.shrikeBT.Compiler.Output;
import com.ibm.wala.shrikeBT.GotoInstruction;
import com.ibm.wala.shrikeBT.IInstruction;
import com.ibm.wala.shrikeBT.MethodData;
import com.ibm.wala.shrikeBT.analysis.Analyzer;
import com.ibm.wala.shrikeBT.analysis.Analyzer.FailureException;
import com.ibm.wala.shrikeBT.analysis.ClassHierarchyProvider;
import com.ibm.wala.shrikeBT.analysis.Verifier;
import com.ibm.wala.shrikeCT.ClassWriter.Element;
import com.ibm.wala.shrikeCT.StackMapConstants.Item;
import com.ibm.wala.shrikeCT.StackMapConstants.ObjectType;
import com.ibm.wala.shrikeCT.StackMapConstants.StackMapFrame;
import com.ibm.wala.shrikeCT.StackMapConstants.StackMapType;
import com.ibm.wala.shrikeCT.StackMapConstants.UninitializedType;
import com.ibm.wala.util.collections.HashMapFactory;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.BitSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

public class StackMapTableWriter extends Element {
  private final byte[] data;

  public StackMapTableWriter(ClassWriter writer, List<StackMapFrame> frames) throws IOException {
    this.data = serialize(writer, frames);
  }

  private static byte[] serialize(ClassWriter writer, List<StackMapFrame> frames)
      throws IOException {
    ByteArrayOutputStream data = new ByteArrayOutputStream();

    for (StackMapFrame frame : frames) {
      frame.write(data, writer);
    }

    ByteArrayOutputStream bs = new ByteArrayOutputStream();
    writeUShort(bs, writer.addCPUtf8("StackMapTable"));
    writeInt(bs, data.size() + 2);
    writeUShort(bs, frames.size());
    data.writeTo(bs);

    return bs.toByteArray();
  }

  public StackMapTableWriter(
      ClassWriter writer,
      MethodData method,
      Output output,
      ClassHierarchyProvider cha,
      String[][] vars)
      throws FailureException, IOException {
    this(writer, stackMapTable(method, output, cha, vars, null));
  }

  public StackMapTableWriter(
      ClassWriter writer,
      MethodData method,
      Output output,
      ClassHierarchyProvider cha,
      String[][] vars,
      List<StackMapFrame> reuseFrames)
      throws FailureException, IOException {
    this(writer, stackMapTable(method, output, cha, vars, reuseFrames));
  }

  private static List<StackMapFrame> remapStackFrames(
      List<StackMapFrame> sm, int[] newBytecodesToOldBytecodes) {
    // mapping to new bytecode
    Map<Integer, Integer> oldToNew = HashMapFactory.make();
    for (int i = newBytecodesToOldBytecodes.length - 1; i >= 0; i--) {
      oldToNew.put(newBytecodesToOldBytecodes[i], i);
    }

    // positions of frames
    int i = 1;
    int positions[] = new int[sm.size()];
    Iterator<StackMapFrame> sms = sm.iterator();
    int position = sms.next().getOffset();
    positions[0] = oldToNew.get(position);
    while (sms.hasNext()) {
      position += sms.next().getOffset() + 1;
      positions[i++] = oldToNew.get(position);
    }

    // positions turned into offsets
    for (i = positions.length - 1; i > 0; i--) {
      positions[i] = positions[i] - positions[i - 1] - 1;
    }

    // frames with new offsets
    List<StackMapFrame> newFrames = new ArrayList<>(sm.size());
    for (i = 0; i < sm.size(); i++) {
      newFrames.add(new StackMapFrame(sm.get(i), positions[i]));
    }

    return newFrames;
  }

  public StackMapTableWriter(
      ClassWriter w, List<StackMapFrame> sm, int[] newBytecodesToOldBytecodes) throws IOException {
    this(w, remapStackFrames(sm, newBytecodesToOldBytecodes));
  }

  static StackMapType item(String type) {
    if (type == null) {
      return Item.ITEM_Top;
    } else if (type.equals(TYPE_null)) {
      return Item.ITEM_Null;
    } else if (type.equals(Analyzer.topType)) {
      return Item.ITEM_Top;
    } else if (type.equals(Analyzer.thisType)) {
      return Item.ITEM_UninitializedThis;
    } else if (type.equals(TYPE_int)) {
      return Item.ITEM_Integer;
    } else if (type.equals(TYPE_float)) {
      return Item.ITEM_Float;
    } else if (type.equals(TYPE_double)) {
      return Item.ITEM_Double;
    } else if (type.equals(TYPE_long)) {
      return Item.ITEM_Long;
    } else {
      if (type.startsWith("#")) {
        return new UninitializedType(type);
      } else {
        return new ObjectType(type);
      }
    }
  }

  static void writeUByte(OutputStream s, int v) throws IOException {
    byte bytes[] = new byte[1];
    ClassWriter.setUByte(bytes, 0, v);
    s.write(bytes);
  }

  static void writeUShort(OutputStream s, int v) throws IOException {
    byte bytes[] = new byte[2];
    ClassWriter.setUShort(bytes, 0, v);
    s.write(bytes);
  }

  static void writeInt(OutputStream s, int v) throws IOException {
    byte bytes[] = new byte[4];
    ClassWriter.setInt(bytes, 0, v);
    s.write(bytes);
  }

  static StackMapType[] trim(StackMapType[] types) {
    int i = types.length - 1;
    while (i >= 0
        && (types[i] == null || types[i] == Item.ITEM_Null || types[i] == Item.ITEM_Top)) {
      i--;
    }

    if (i < 0) {
      return new StackMapType[0];
    } else if (i < types.length - 1) {
      StackMapType[] trimmed = new StackMapType[i + 1];
      System.arraycopy(types, 0, trimmed, 0, i + 1);
      return trimmed;
    } else {
      return types;
    }
  }

  private static String hackUnknown(String type) {
    if (type == null) {
      return type;
    } else if (type.startsWith("[")) {
      return '[' + hackUnknown(type.substring(1));
    } else if ("L?;".equals(type)) {
      return "Ljava/lang/Object;";
    } else {
      return type;
    }
  }

  static StackMapType[] types(String[] types, boolean locals) {
    StackMapType[] stackTypes = new StackMapType[types.length];

    int x = 0;
    for (int j = 0; j < types.length; j++) {
      StackMapType stackType = item(hackUnknown(types[j]));
      stackTypes[x++] = stackType;
      if (locals && stackType.size() == 2) {
        j++;
      }
    }

    return trim(stackTypes);
  }

  private static boolean isUselessGoto(IInstruction inst, int index) {
    if (inst instanceof GotoInstruction) {
      if (((GotoInstruction) inst).getBranchTargets()[0] == index + 1) {
        return true;
      }
    }

    return false;
  }

  public static List<StackMapFrame> stackMapTable(
      MethodData method,
      Output output,
      ClassHierarchyProvider cha,
      String[][] vars,
      List<StackMapFrame> reuseFrames)
      throws FailureException {
    int idx = 0;

    List<StackMapFrame> frames = new ArrayList<>();

    int[] instructionToBytecode = output.getInstructionOffsets();
    IInstruction[] insts = method.getInstructions();

    Verifier typeChecker = new Verifier(method, instructionToBytecode, vars);
    if (cha != null) {
      typeChecker.setClassHierarchy(cha);
    }
    typeChecker.computeTypes();
    BitSet bbs = typeChecker.getBasicBlockStarts();

    int offset = 0;
    for (int i = 1; i < insts.length; i++) {
      if (bbs.get(i)) {

        // Shrike does not generate goto i+1
        if (isUselessGoto(insts[i], i)) {
          continue;
        }

        // offset delta
        int position = instructionToBytecode[i];
        assert position - offset > 0 || offset == 0;
        int frameOffset = offset == 0 ? position : position - offset - 1;
        offset = position;

        if (reuseFrames != null) {
          if (idx < reuseFrames.size() && reuseFrames.get(idx).getOffset() == frameOffset) {
            frames.add(reuseFrames.get(idx++));
            continue;
          } else {
            reuseFrames = null;
          }
        }

        // full frame
        byte frameType = (byte) 255;

        // locals
        String[] localTypes = typeChecker.getLocalTypes()[i];
        StackMapType[] localWriteTypes;
        if (localTypes != null) {
          localWriteTypes = types(localTypes, true);
        } else {
          localWriteTypes = new StackMapType[0];
        }

        // stack
        String[] stackTypes = typeChecker.getStackTypes()[i];
        StackMapType[] stackWriteTypes;
        if (stackTypes != null) {
          stackWriteTypes = types(stackTypes, false);
        } else {
          stackWriteTypes = new StackMapType[0];
        }

        frames.add(new StackMapFrame(frameType, frameOffset, localWriteTypes, stackWriteTypes));
      }
    }

    return frames;
  }

  @Override
  public int getSize() {
    return data.length;
  }

  @Override
  public int copyInto(byte[] buf, int offset) {
    System.arraycopy(data, 0, buf, offset + 0, data.length);
    return data.length + offset;
  }
}
