/*
 * Copyright 2011 Henry Coles
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and limitations under the License.
 */
package org.pitest.coverage.execute;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.pitest.bytecode.FrameOptions;
import org.pitest.classinfo.ComputeClassWriter;
import org.pitest.classpath.ClassloaderByteArraySource;
import org.pitest.coverage.AlreadyInstrumentedException;
import org.pitest.coverage.CoverageClassVisitor;
import org.pitest.reflection.Reflection;
import org.pitest.util.IsolationUtils;
import org.pitest.util.StreamUtil;
import org.pitest.util.Unchecked;
import sun.pitest.CodeCoverageStore;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public final class JavassistCoverageInterceptor {

  private static final Map<String, String> COMPUTE_CACHE = new ConcurrentHashMap<>();

  private JavassistCoverageInterceptor() {

  }

  public static InputStream openClassfile(final Object classPath, // NO_UCD
      final String name) {

    try {
      if (isInstrumentedClass(name)) {
        final byte[] bs = getOriginalBytes(classPath, name);
        return new ByteArrayInputStream(
            transformBytes(IsolationUtils.getContextClassLoader(), name, bs));
      } else {
        return returnNormalBytes(classPath, name);
      }
    } catch (final IOException ex) {
      throw Unchecked.translateCheckedException(ex);
    }

  }

  private static byte[] getOriginalBytes(final Object classPath,
      final String name) throws IOException {
    try (InputStream is = returnNormalBytes(classPath,name)) {
      return StreamUtil.streamToByteArray(is);
    }
  }

  private static byte[] transformBytes(final ClassLoader loader,
      final String className, final byte[] classfileBuffer) {
    final ClassReader reader = new ClassReader(classfileBuffer);
    final ClassWriter writer = new ComputeClassWriter(
        new ClassloaderByteArraySource(loader), COMPUTE_CACHE,
        FrameOptions.pickFlags(classfileBuffer));

    // The transformed classes will be given a different id than the one already loaded.
    // Not clear if this is desirable or not. At the point of writing this comment
    // pitest will merge coverage of all classes with the same fully qualified name.
    // If this changes this might become a bug, however it would also probably not be possible
    // to support powermock if this assumption changed, so this code would most likely be deleted.

    final int id = CodeCoverageStore.registerClass(className);
    try {
      reader.accept(new CoverageClassVisitor(id, writer),
          ClassReader.EXPAND_FRAMES);
      return writer.toByteArray();
    } catch (AlreadyInstrumentedException ex) {
      return null;
    }
  }

  private static InputStream returnNormalBytes(final Object classPath,
      final String name) {
    try {
      return (InputStream) Reflection.publicMethod(classPath.getClass(),
          "openClassfile").invoke(classPath, name);
    } catch (final Exception e) {
      throw Unchecked.translateCheckedException(e);
    }
  }

  private static boolean isInstrumentedClass(final String name) {
    return true;
  }


}
