/*
 * Copyright (c) 2002,2006 IBM Corporation.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * which accompanies this distribution, and is available at
 * http://www.eclipse.org/legal/epl-v10.html
 *
 * Contributors:
 *     IBM Corporation - initial API and implementation
 */
package com.ibm.wala.shrike.copywriter;

import com.ibm.wala.shrikeBT.Compiler;
import com.ibm.wala.shrikeBT.Decoder.InvalidBytecodeException;
import com.ibm.wala.shrikeBT.MethodData;
import com.ibm.wala.shrikeBT.shrikeCT.CTCompiler;
import com.ibm.wala.shrikeBT.shrikeCT.CTDecoder;
import com.ibm.wala.shrikeBT.shrikeCT.ClassInstrumenter;
import com.ibm.wala.shrikeBT.shrikeCT.OfflineInstrumenter;
import com.ibm.wala.shrikeCT.ClassConstants;
import com.ibm.wala.shrikeCT.ClassReader;
import com.ibm.wala.shrikeCT.ClassReader.AttrIterator;
import com.ibm.wala.shrikeCT.ClassWriter;
import com.ibm.wala.shrikeCT.ClassWriter.Element;
import com.ibm.wala.shrikeCT.CodeReader;
import com.ibm.wala.shrikeCT.CodeWriter;
import com.ibm.wala.shrikeCT.ConstantPoolParser;
import com.ibm.wala.shrikeCT.ConstantValueReader;
import com.ibm.wala.shrikeCT.ConstantValueWriter;
import com.ibm.wala.shrikeCT.ExceptionsReader;
import com.ibm.wala.shrikeCT.ExceptionsWriter;
import com.ibm.wala.shrikeCT.InnerClassesReader;
import com.ibm.wala.shrikeCT.InnerClassesWriter;
import com.ibm.wala.shrikeCT.InvalidClassFileException;
import com.ibm.wala.shrikeCT.LocalVariableTableReader;
import com.ibm.wala.shrikeCT.LocalVariableTableWriter;
import com.ibm.wala.shrikeCT.SourceFileReader;
import com.ibm.wala.shrikeCT.SourceFileWriter;
import java.io.OutputStreamWriter;
import java.io.Writer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.zip.ZipEntry;

public class CopyWriter {
  private static final String USAGE =
      "IBM CopyWriter Tool\n"
          + "This tool takes the following command line options:\n"
          + "    <jarname> <jarname> ...   Process the classes from these jars\n"
          + "    -o <jarname>              Put the resulting classes into <jarname>\n"
          + "    -c <copyright>            Make the copyright string be\n"
          + "                              '\u00A9 Copyright <copyright>'";

  private static OfflineInstrumenter instrumenter;

  public static String copyright;

  public static final String copyrightAttrName = "com.ibm.Copyright";

  private int replaceWith;

  private int replace;

  static class UnknownAttributeException extends Exception {
    private static final long serialVersionUID = 8845177787110364793L;

    UnknownAttributeException(String t) {
      super("Attribute '" + t + "' not understood");
    }
  }

  public static void main(String[] args) throws Exception {
    if (args == null || args.length == 0) {
      System.err.println(USAGE);
      System.exit(1);
    }

    for (int i = 0; i < args.length - 1; i++) {
      if (args[i] == null) {
        throw new IllegalArgumentException("args[" + i + "] is null");
      }
      if (args[i].equals("-c")) {
        copyright = "\u00A9 Copyright " + args[i + 1];
        String[] newArgs = new String[args.length - 2];
        System.arraycopy(args, 0, newArgs, 0, i);
        System.arraycopy(args, i + 2, newArgs, i, newArgs.length - i);
        args = newArgs;
        break;
      }
    }
    if (copyright == null) {
      System.err.println(USAGE);
      System.exit(1);
    }

    final ArrayList<ZipEntry> entries = new ArrayList<>();

    instrumenter = new OfflineInstrumenter();
    instrumenter.setManifestBuilder(entries::add);
    instrumenter.parseStandardArgs(args);
    instrumenter.setJARComment(copyright);
    instrumenter.beginTraversal();
    ClassInstrumenter ci;
    CopyWriter cw = new CopyWriter();
    while ((ci = instrumenter.nextClass()) != null) {
      try {
        cw.doClass(ci);
      } catch (UnknownAttributeException ex) {
        System.err.println(ex.getMessage() + " in " + instrumenter.getLastClassResourceName());
      }
    }

    instrumenter.writeUnmodifiedClasses();

    Writer w =
        new OutputStreamWriter(instrumenter.addOutputJarEntry(new ZipEntry("IBM-Copyright")));
    w.write(copyright + '\n');
    for (ZipEntry ze : entries) {
      w.write("  " + ze.getName() + '\n');
    }
    w.write(copyright + '\n');
    w.flush();
    instrumenter.endOutputJarEntry();

    instrumenter.close();
  }

  private int transformCPIndex(int i) {
    if (i == replace) {
      return replaceWith;
    } else {
      return i;
    }
  }

  private Element transformAttribute(ClassReader cr, int m, ClassWriter w, AttrIterator iter)
      throws InvalidClassFileException, UnknownAttributeException, InvalidBytecodeException {
    String name = iter.getName();

    boolean needTransform = true;
    if (name.equals("Synthetic") || name.equals("Deprecated") || name.equals("LineNumberTable")) {
      needTransform = false;
    }

    int offset = iter.getRawOffset();
    int end = offset + iter.getRawSize();

    if (needTransform) {
      needTransform = false;
      for (int i = offset; i + 1 < end; i++) {
        if (cr.getUShort(i) == replace) {
          break;
        }
      }
    }

    if (!needTransform) {
      return new ClassWriter.RawElement(cr.getBytes(), offset, end - offset);
    }

    switch (name) {
      case "Code":
        {
          CodeReader r = new CodeReader(iter);
          CTDecoder decoder = new CTDecoder(r);
          decoder.decode();
          MethodData md =
              new MethodData(
                  decoder,
                  cr.getMethodAccessFlags(m),
                  CTDecoder.convertClassToType(cr.getName()),
                  cr.getMethodName(m),
                  cr.getMethodType(m));
          CTCompiler compiler = CTCompiler.make(w, md);
          compiler.compile();
          if (compiler.getAuxiliaryMethods().length > 0)
            throw new Error("Where did this auxiliary method come from?");
          Compiler.Output out = compiler.getOutput();
          CodeWriter cw = new CodeWriter(w);
          cw.setMaxLocals(out.getMaxLocals());
          cw.setMaxStack(out.getMaxStack());
          cw.setCode(out.getCode());
          cw.setRawHandlers(out.getRawHandlers());
          AttrIterator iterator = new AttrIterator();
          r.initAttributeIterator(iterator);
          cw.setAttributes(collectAttributes(cr, m, w, iterator));
          return cw;
        }
      case "ConstantValue":
        {
          ConstantValueReader r = new ConstantValueReader(iter);
          ConstantValueWriter cw = new ConstantValueWriter(w);
          cw.setValueCPIndex(transformCPIndex(r.getValueCPIndex()));
          return cw;
        }
      case "SourceFile":
        {
          SourceFileReader r = new SourceFileReader(iter);
          SourceFileWriter cw = new SourceFileWriter(w);
          cw.setSourceFileCPIndex(transformCPIndex(r.getSourceFileCPIndex()));
          return cw;
        }
      case "LocalVariableTableReader":
        {
          LocalVariableTableReader lr = new LocalVariableTableReader(iter);
          LocalVariableTableWriter lw = new LocalVariableTableWriter(w);
          int[] table = lr.getRawTable();
          for (int i = 0; i < table.length; i += 5) {
            table[i + 2] = transformCPIndex(table[i + 2]);
            table[i + 3] = transformCPIndex(table[i + 3]);
          }
          lw.setRawTable(table);
          return lw;
        }
      case "Exceptions":
        {
          ExceptionsReader lr = new ExceptionsReader(iter);
          ExceptionsWriter lw = new ExceptionsWriter(w);
          int[] table = lr.getRawTable();
          Arrays.setAll(table, i -> transformCPIndex(table[i]));
          lw.setRawTable(table);
          return lw;
        }
      case "InnerClasses":
        {
          InnerClassesReader lr = new InnerClassesReader(iter);
          InnerClassesWriter lw = new InnerClassesWriter(w);
          int[] table = lr.getRawTable();
          for (int i = 0; i < table.length; i += 4) {
            table[i] = transformCPIndex(table[i]);
            table[i + 1] = transformCPIndex(table[i + 1]);
            table[i + 2] = transformCPIndex(table[i + 2]);
          }
          lw.setRawTable(table);
          return lw;
        }
    }

    throw new UnknownAttributeException(name);
  }

  private Element[] collectAttributes(ClassReader cr, int m, ClassWriter w, AttrIterator iter)
      throws InvalidClassFileException, UnknownAttributeException, InvalidBytecodeException {
    Element[] elems = new Element[iter.getRemainingAttributesCount()];
    for (int i = 0; i < elems.length; i++) {
      elems[i] = transformAttribute(cr, m, w, iter);
      iter.advance();
    }
    return elems;
  }

  private static int copyEntry(ConstantPoolParser cp, ClassWriter w, int i)
      throws InvalidClassFileException {
    byte t = cp.getItemType(i);
    switch (t) {
      case ClassConstants.CONSTANT_String:
        return w.addCPString(cp.getCPString(i));
      case ClassConstants.CONSTANT_Class:
        return w.addCPClass(cp.getCPClass(i));
      case ClassConstants.CONSTANT_FieldRef:
        return w.addCPFieldRef(cp.getCPRefClass(i), cp.getCPRefName(i), cp.getCPRefType(i));
      case ClassConstants.CONSTANT_InterfaceMethodRef:
        return w.addCPInterfaceMethodRef(
            cp.getCPRefClass(i), cp.getCPRefName(i), cp.getCPRefType(i));
      case ClassConstants.CONSTANT_MethodRef:
        return w.addCPMethodRef(cp.getCPRefClass(i), cp.getCPRefName(i), cp.getCPRefType(i));
      case ClassConstants.CONSTANT_NameAndType:
        return w.addCPNAT(cp.getCPNATName(i), cp.getCPNATType(i));
      case ClassConstants.CONSTANT_Integer:
        return w.addCPInt(cp.getCPInt(i));
      case ClassConstants.CONSTANT_Float:
        return w.addCPFloat(cp.getCPFloat(i));
      case ClassConstants.CONSTANT_Long:
        return w.addCPLong(cp.getCPLong(i));
      case ClassConstants.CONSTANT_Double:
        return w.addCPDouble(cp.getCPDouble(i));
      case ClassConstants.CONSTANT_Utf8:
        return w.addCPUtf8(cp.getCPUtf8(i));
      default:
        return -1;
    }
  }

  private void doClass(final ClassInstrumenter ci) throws Exception {
    /*
     * Our basic strategy is to make the first element of the constant pool be the copyright string (as a UTF8 constant pool item).
     * This requires us to parse and emit any class data which might refer to that constant pool item (#1). We will assume that any
     * attribute which refers to that constant pool item must contain the byte sequence '00 01', so we can just copy over any
     * attributes which don't contain that byte sequence. If we detect an unknown attribute type containing the sequence '00 01',
     * then we will abort.
     */
    ClassReader cr = ci.getReader();
    ClassWriter w = new ClassWriter();

    // Make sure that when we're moving over the old constant pool,
    // anytime we add a new entry it really is added and we don't just
    // reuse an existing entry.
    w.setForceAddCPEntries(true);

    // Make the first string in the constant pool be the copyright string
    int r = w.addCPUtf8(copyright);
    if (r != 1) throw new Error("Invalid constant pool index: " + r);

    // Now add the rest of the CP entries
    ConstantPoolParser cp = cr.getCP();
    int CPCount = cp.getItemCount();

    if (1 < CPCount) {
      final byte itemType = cp.getItemType(1);
      switch (itemType) {
        case ClassConstants.CONSTANT_Long:
        case ClassConstants.CONSTANT_Double:
          // item 1 is a double-word item, so the next real item is at 3
          // to make sure item 3 is allocated at index 3, we'll need to
          // insert a dummy entry at index 2
          r = w.addCPUtf8("");
          if (r != 2) throw new Error("Invalid constant pool index for dummy: " + r);
          break;
        default:
          throw new UnsupportedOperationException(
              String.format("unexpected constant-pool item type %s", itemType));
      }
    }
    for (int i = 2; i < CPCount; i++) {
      r = copyEntry(cp, w, i);
      if (r != -1 && r != i)
        throw new Error("Invalid constant pool index allocated: " + r + ", expected " + i);
    }
    w.setForceAddCPEntries(false);

    // add CP entry we replaced
    replaceWith = copyEntry(cp, w, 1);
    replace = 1;

    // emit class
    w.setMajorVersion(cr.getMajorVersion());
    w.setMinorVersion(cr.getMinorVersion());
    w.setAccessFlags(cr.getAccessFlags());
    w.setName(cr.getName());
    w.setSuperName(cr.getSuperName());
    w.setInterfaceNames(cr.getInterfaceNames());

    ClassReader.AttrIterator iter = new ClassReader.AttrIterator();

    int fieldCount = cr.getFieldCount();
    for (int i = 0; i < fieldCount; i++) {
      cr.initFieldAttributeIterator(i, iter);
      w.addField(
          cr.getFieldAccessFlags(i),
          cr.getFieldName(i),
          cr.getFieldType(i),
          collectAttributes(cr, i, w, iter));
    }

    int methodCount = cr.getMethodCount();
    for (int i = 0; i < methodCount; i++) {
      cr.initMethodAttributeIterator(i, iter);
      w.addMethod(
          cr.getMethodAccessFlags(i),
          cr.getMethodName(i),
          cr.getMethodType(i),
          collectAttributes(cr, i, w, iter));
    }

    cr.initClassAttributeIterator(iter);
    for (; iter.isValid(); iter.advance()) {
      w.addClassAttribute(transformAttribute(cr, 0, w, iter));
    }

    instrumenter.outputModifiedClass(ci, w);
  }
}
