/*
 * #%L
 * Betamax Arquillian Wildfly :: Extension
 * %%
 * Copyright (C) 2016 James Netherton
 * %%
 * 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.
 * #L%
 */
package com.github.jamesnetherton.betamax.arquillian;

import com.github.jamesnetherton.betamax.arquillian.utils.ManifestBuilder;
import org.jboss.arquillian.container.test.spi.TestDeployment;
import org.jboss.arquillian.container.test.spi.client.deployment.ProtocolArchiveProcessor;
import org.jboss.shrinkwrap.api.Archive;
import org.jboss.shrinkwrap.api.ArchivePath;
import org.jboss.shrinkwrap.api.ArchivePaths;
import org.jboss.shrinkwrap.api.Node;
import org.jboss.shrinkwrap.api.asset.Asset;
import org.jboss.shrinkwrap.api.asset.ClassAsset;
import software.betamax.junit.Betamax;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.lang.reflect.Method;
import java.util.Map;
import java.util.jar.Attributes;
import java.util.jar.JarFile;
import java.util.jar.Manifest;

/**
 * A {@link ProtocolArchiveProcessor} implementation that adds Betamax module dependencies to the deployment MANIFEST.MF
 */
public class BetamaxArchiveProcessor implements ProtocolArchiveProcessor {
    private static final String[] BETAMAX_DEPENDENCIES = {
        "com.google.guava",
        "org.apache.commons.cli",
        "org.apache.commons.codec",
        "org.slf4j",
        "org.yaml.snakeyaml",
        "io.netty"
    };

    @Override
    public void process(TestDeployment testDeployment, Archive<?> protocolArchive) {
        Archive<?> applicationArchive = testDeployment.getApplicationArchive();
        if (!isValidArchive(applicationArchive)) {
            return;
        }

        // Add Betamax dependencies to deployment manifest
        Manifest manifest = getManifest(protocolArchive);
        String betamaxDependencies = new ManifestBuilder().addDepencencies(BETAMAX_DEPENDENCIES).build();
        String dependencies = manifest.getMainAttributes().getValue("Dependencies");
        manifest.getMainAttributes().putValue("Dependencies", dependencies + betamaxDependencies);

        ArchivePath manifestPath = ArchivePaths.create(JarFile.MANIFEST_NAME);
        protocolArchive.delete(manifestPath);
        protocolArchive.add((Asset) () -> {
            try {
                ByteArrayOutputStream baos = new ByteArrayOutputStream();
                manifest.write(baos);
                return new ByteArrayInputStream(baos.toByteArray());
            } catch (IOException ex) {
                throw new IllegalStateException("Failed writing manifest", ex);
            }
        }, manifestPath);
    }

    private Manifest getManifest(Archive<?> archive) {
        Node node = archive.get(JarFile.MANIFEST_NAME);
        Manifest manifest;
        try {
            if (node == null) {
                manifest = new Manifest();
                Attributes attributes = manifest.getMainAttributes();
                attributes.putValue(Attributes.Name.MANIFEST_VERSION.toString(), "1.0");
            } else {
                manifest = new Manifest(node.getAsset().openStream());
            }
        } catch (IOException e) {
            throw new IllegalStateException("Failed to obtain deployment manifest", e);
        }
        return manifest;
    }

    private boolean isValidArchive(Archive<?> archive) {
        Map<ArchivePath, Node> content = archive.getContent();

        for (ArchivePath path : content.keySet()) {
            Asset asset = content.get(path).getAsset();
            if (asset instanceof ClassAsset) {
                Class<?> clazz = ((ClassAsset) asset).getSource();

                for (Method method : clazz.getMethods()) {
                    Betamax betamax = BetamaxAnnotationFinder.find(method);
                    if (betamax != null) {
                        return true;
                    }
                }
            }
        }

        return false;
    }
}
