package hudson.plugins.ec2.one;

import hudson.Extension;
import hudson.Util;
import hudson.model.Describable;
import hudson.model.TaskListener;
import hudson.model.Descriptor;
import hudson.model.Descriptor.FormException;
import hudson.model.Hudson;
import hudson.model.Label;
import hudson.model.Node;
import hudson.model.labels.LabelAtom;
import hudson.plugins.ec2.EC2Computer;
import hudson.plugins.ec2.EC2Slave;
import hudson.util.FormValidation;

import java.io.IOException;
import java.io.PrintStream;
import java.net.URL;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;

import javax.servlet.ServletException;

import org.apache.commons.codec.binary.Base64;
import org.apache.commons.lang.StringUtils;
import org.kohsuke.stapler.DataBoundConstructor;
import org.kohsuke.stapler.QueryParameter;

import com.amazonaws.AmazonClientException;
import com.amazonaws.services.ec2.AmazonEC2;
import com.amazonaws.services.ec2.model.CreateTagsRequest;
import com.amazonaws.services.ec2.model.DescribeImagesRequest;
import com.amazonaws.services.ec2.model.DescribeInstancesRequest;
import com.amazonaws.services.ec2.model.DescribeInstancesResult;
import com.amazonaws.services.ec2.model.DescribeSecurityGroupsRequest;
import com.amazonaws.services.ec2.model.DescribeSecurityGroupsResult;
import com.amazonaws.services.ec2.model.DescribeSubnetsRequest;
import com.amazonaws.services.ec2.model.DescribeSubnetsResult;
import com.amazonaws.services.ec2.model.Filter;
import com.amazonaws.services.ec2.model.Image;
import com.amazonaws.services.ec2.model.Instance;
import com.amazonaws.services.ec2.model.InstanceStateName;
import com.amazonaws.services.ec2.model.KeyPair;
import com.amazonaws.services.ec2.model.Placement;
import com.amazonaws.services.ec2.model.Reservation;
import com.amazonaws.services.ec2.model.RunInstancesRequest;
import com.amazonaws.services.ec2.model.SecurityGroup;
import com.amazonaws.services.ec2.model.StartInstancesRequest;
import com.amazonaws.services.ec2.model.StartInstancesResult;
import com.amazonaws.services.ec2.model.Tag;

/**
 * Template of {@link EC2Slave} to launch.
 * 
 * @author Kohsuke Kawaguchi
 */
public class OneSlaveTemplate implements Describable<OneSlaveTemplate> {
    public final String ami;

    public final String description;

    public final String zone;

    public final String securityGroups;

    public final String remoteFS;

    public final String sshPort;

    public final OneInstanceType type;

    public final String labels;

    public final String initScript;

    public final String userData;

    public final String numExecutors;

    public final String remoteAdmin;

    public final String rootCommandPrefix;

    public final String jvmopts;

    public final String subnetId;

    public final String idleTerminationMinutes;

    public final int instanceCap;

    public final boolean stopOnTerminate;

    private final List<OneEC2Tag> tags;

    public final boolean usePrivateDnsName;

    protected transient OneEC2Cloud parent;

    private transient/* almost final */Set<LabelAtom> labelSet;

    private transient/* almost final */Set<String> securityGroupSet;

    @DataBoundConstructor
    public OneSlaveTemplate(final String ami, final String zone,
            final String securityGroups, final String remoteFS,
            final String sshPort, final OneInstanceType type,
            final String labelString, final String description,
            final String initScript, final String userData,
            final String numExecutors, final String remoteAdmin,
            final String rootCommandPrefix, final String jvmopts,
            final boolean stopOnTerminate, final String subnetId,
            final List<OneEC2Tag> tags,
            final String idleTerminationMinutes,
            final boolean usePrivateDnsName, final String instanceCapStr) {
        this.ami = ami;
        this.zone = zone;
        this.securityGroups = securityGroups;
        this.remoteFS = remoteFS;
        this.sshPort = sshPort;
        this.type = type;
        labels = Util.fixNull(labelString);
        this.description = description;
        this.initScript = initScript;
        this.userData = userData;
        this.numExecutors = Util.fixNull(numExecutors).trim();
        this.remoteAdmin = remoteAdmin;
        this.rootCommandPrefix = rootCommandPrefix;
        this.jvmopts = jvmopts;
        this.stopOnTerminate = stopOnTerminate;
        this.subnetId = subnetId;
        this.tags = tags;
        this.idleTerminationMinutes = idleTerminationMinutes;
        this.usePrivateDnsName = usePrivateDnsName;

        if (null == instanceCapStr || instanceCapStr.equals("")) {
            instanceCap = Integer.MAX_VALUE;
        } else {
            instanceCap = Integer.parseInt(instanceCapStr);
        }

        readResolve(); // initialize
    }

    public OneEC2Cloud getParent() {
        return parent;
    }

    public String getLabelString() {
        return labels;
    }

    public String getDisplayName() {
        return description + " (" + ami + ")";
    }

    String getZone() {
        return zone;
    }

    public String getSecurityGroupString() {
        return securityGroups;
    }

    public Set<String> getSecurityGroupSet() {
        return securityGroupSet;
    }

    public Set<String> parseSecurityGroups() {
        if (securityGroups == null || "".equals(securityGroups.trim())) {
            return Collections.emptySet();
        } else {
            return new HashSet<String>(Arrays.asList(securityGroups
                .split("\\s*,\\s*")));
        }
    }

    public int getNumExecutors() {
        try {
            return Integer.parseInt(numExecutors);
        } catch (NumberFormatException e) {
            return OneEC2Slave.toNumExecutors(type);
        }
    }

    public int getSshPort() {
        try {
            return Integer.parseInt(sshPort);
        } catch (NumberFormatException e) {
            return 22;
        }
    }

    public String getRemoteAdmin() {
        return remoteAdmin;
    }

    public String getRootCommandPrefix() {
        return rootCommandPrefix;
    }

    public String getSubnetId() {
        return subnetId;
    }

    public List<OneEC2Tag> getTags() {
        if (null == tags) {
            return null;
        }
        return Collections.unmodifiableList(tags);
    }

    public String getidleTerminationMinutes() {
        return idleTerminationMinutes;
    }

    public Set<LabelAtom> getLabelSet() {
        return labelSet;
    }

    public int getInstanceCap() {
        return instanceCap;
    }

    public String getInstanceCapStr() {
        if (instanceCap == Integer.MAX_VALUE) {
            return "";
        } else {
            return String.valueOf(instanceCap);
        }
    }

    /**
     * Does this contain the given label?
     * 
     * @param l
     *            can be null to indicate "don't care".
     */
    public boolean containsLabel(final Label l) {
        return l == null || labelSet.contains(l);
    }

    /**
     * Provisions a new EC2 slave.
     * 
     * @return always non-null. This needs to be then added to
     *         {@link Hudson#addNode(Node)}.
     */
    public OneEC2Slave provision(final TaskListener listener)
            throws AmazonClientException, IOException {
        PrintStream logger = listener.getLogger();
        AmazonEC2 ec2 = getParent().connect();

        try {
            logger.println("Launching " + ami);
            KeyPair keyPair =
                (null != parent.getPrivateKey()) ? parent.getPrivateKey()
                    .find(ec2) : null;

            RunInstancesRequest riRequest =
                new RunInstancesRequest(ami, 1, 1);

            List<Filter> diFilters = new ArrayList<Filter>();
            diFilters.add(new Filter("image-id").withValues(ami));

            if (StringUtils.isNotBlank(getZone())) {
                Placement placement = new Placement(getZone());
                riRequest.setPlacement(placement);
                diFilters.add(new Filter("availability-zone")
                    .withValues(getZone()));
            }

            if (StringUtils.isNotBlank(getSubnetId())) {
                riRequest.setSubnetId(getSubnetId());
                diFilters.add(new Filter("subnet-id")
                    .withValues(getSubnetId()));

                /*
                 * If we have a subnet ID then we can only use VPC security
                 * groups
                 */
                if (!securityGroupSet.isEmpty()) {
                    List<String> group_ids = new ArrayList<String>();

                    DescribeSecurityGroupsRequest group_req =
                        new DescribeSecurityGroupsRequest();
                    group_req.withFilters(new Filter("group-name")
                        .withValues(securityGroupSet));
                    DescribeSecurityGroupsResult group_result =
                        ec2.describeSecurityGroups(group_req);

                    for (SecurityGroup group : group_result
                        .getSecurityGroups()) {
                        if (group.getVpcId() != null
                            && !group.getVpcId().isEmpty()) {
                            List<Filter> filters = new ArrayList<Filter>();
                            filters.add(new Filter("vpc-id")
                                .withValues(group.getVpcId()));
                            filters.add(new Filter("state")
                                .withValues("available"));
                            filters.add(new Filter("subnet-id")
                                .withValues(getSubnetId()));

                            DescribeSubnetsRequest subnet_req =
                                new DescribeSubnetsRequest();
                            subnet_req.withFilters(filters);
                            DescribeSubnetsResult subnet_result =
                                ec2.describeSubnets(subnet_req);

                            List subnets = subnet_result.getSubnets();
                            if (subnets != null && !subnets.isEmpty()) {
                                group_ids.add(group.getGroupId());
                            }
                        }
                    }

                    if (securityGroupSet.size() != group_ids.size()) {
                        throw new AmazonClientException(
                            "Security groups must all be VPC security groups to work in a VPC context");
                    }

                    if (!group_ids.isEmpty()) {
                        riRequest.setSecurityGroupIds(group_ids);
                        diFilters.add(new Filter("instance.group-id")
                            .withValues(group_ids));
                    }
                }
            } else {
                /* No subnet: we can use standard security groups by name */
                riRequest.setSecurityGroups(securityGroupSet);
                if (securityGroupSet.size() > 0) {
                    diFilters.add(new Filter("group-name")
                        .withValues(securityGroupSet));
                }
            }

            if (null != userData) {
                String userDataString =
                    Base64.encodeBase64String(userData.getBytes());
                riRequest.setUserData(userDataString);
            }
            if (null != keyPair) {
                riRequest.setKeyName(keyPair.getKeyName());
                diFilters.add(new Filter("key-name").withValues(keyPair
                    .getKeyName()));
            }
            riRequest.setInstanceType(type.toString());
            diFilters.add(new Filter("instance-type").withValues(type
                .toString()));

            HashSet<Tag> inst_tags = null;
            if (tags != null && !tags.isEmpty()) {
                inst_tags = new HashSet<Tag>();
                for (OneEC2Tag t : tags) {
                    diFilters.add(new Filter("tag:" + t.getName())
                        .withValues(t.getValue()));
                }
            }

            DescribeInstancesRequest diRequest =
                new DescribeInstancesRequest();
            diFilters.add(new Filter("instance-state-name").withValues(
                InstanceStateName.Stopped.toString(),
                InstanceStateName.Stopping.toString()));
            diRequest.setFilters(diFilters);
            logger.println("Looking for existing instances: " + diRequest);

            DescribeInstancesResult diResult =
                ec2.describeInstances(diRequest);

            filterReservations(diResult.getReservations(), diFilters);

            if (diResult.getReservations().size() == 0) {
                // Have to create a new instance
                Instance inst =
                    ec2.runInstances(riRequest).getReservation()
                        .getInstances().get(0);

                /* Now that we have our instance, we can set tags on it */
                if (inst_tags != null) {
                    CreateTagsRequest tag_request =
                        new CreateTagsRequest();
                    tag_request.withResources(inst.getInstanceId())
                        .setTags(inst_tags);
                    ec2.createTags(tag_request);

                    // That was a remote request - we should also update our
                    // local instance data.
                    inst.setTags(inst_tags);
                }
                logger.println("No existing instance found - created: "
                    + inst);
                return newSlave(inst);
            }

            Instance inst =
                diResult.getReservations().get(0).getInstances().get(0);
            logger.println("Found existing stopped instance: " + inst);
            List<String> instances = new ArrayList<String>();
            instances.add(inst.getInstanceId());
            StartInstancesRequest siRequest =
                new StartInstancesRequest(instances);
            StartInstancesResult siResult = ec2.startInstances(siRequest);
            logger.println("Starting existing instance: " + inst
                + " result:" + siResult);

            List<Node> nodes = Hudson.getInstance().getNodes();
            for (int i = 0, len = nodes.size(); i < len; i++) {
                if (!(nodes.get(i) instanceof EC2Slave)) {
                    continue;
                }
                OneEC2Slave ec2Node = (OneEC2Slave) nodes.get(i);
                if (ec2Node.getInstanceId().equals(inst.getInstanceId())) {
                    logger.println("Found existing corresponding: "
                        + ec2Node);
                    return ec2Node;
                }
            }

            // Existing slave not found
            logger.println("Creating new slave for existing instance: "
                + inst);
            return newSlave(inst);

        } catch (FormException e) {
            throw new AssertionError(); // we should have discovered all
                                        // configuration issues upfront
        }
    }

    private OneEC2Slave newSlave(final Instance inst)
            throws FormException, IOException {
        return new OneEC2Slave(inst.getInstanceId(), description,
            remoteFS, getSshPort(), getNumExecutors(), labels, initScript,
            remoteAdmin, rootCommandPrefix, jvmopts, stopOnTerminate,
            idleTerminationMinutes, inst.getPublicDnsName(),
            inst.getPrivateDnsName(), OneEC2Tag.fromAmazonTags(inst
                .getTags()), usePrivateDnsName);
    }

    /**
     * Provisions a new EC2 slave based on the currently running instance on
     * EC2, instead of starting a new one.
     */
    public OneEC2Slave attach(final String instanceId,
            final TaskListener listener)
            throws AmazonClientException, IOException {
        PrintStream logger = listener.getLogger();
        AmazonEC2 ec2 = getParent().connect();

        try {
            logger.println("Attaching to " + instanceId);
            DescribeInstancesRequest request =
                new DescribeInstancesRequest();
            request.setInstanceIds(Collections.singletonList(instanceId));
            Instance inst =
                ec2.describeInstances(request).getReservations().get(0)
                    .getInstances().get(0);
            return newSlave(inst);
        } catch (FormException e) {
            throw new AssertionError(); // we should have discovered all
                                        // configuration issues upfront
        }
    }

    /**
     * Initializes data structure that we don't persist.
     */
    protected Object readResolve() {
        labelSet = Label.parse(labels);
        securityGroupSet = parseSecurityGroups();
        return this;
    }

    public Descriptor<OneSlaveTemplate> getDescriptor() {
        return Hudson.getInstance().getDescriptor(getClass());
    }

    @Extension
    public static final class DescriptorImpl extends
            Descriptor<OneSlaveTemplate> {
        @Override
        public String getDisplayName() {
            return null;
        }

        /**
         * Since this shares much of the configuration with {@link EC2Computer},
         * check its help page, too.
         */
        @Override
        public String getHelpFile(final String fieldName) {
            String p = super.getHelpFile(fieldName);
            if (p == null) {
                p =
                    Hudson.getInstance().getDescriptor(OneEC2Slave.class)
                        .getHelpFile(fieldName);
            }
            return p;
        }

        /***
         * Check that the AMI requested is available in the cloud and can be
         * used.
         */
        public FormValidation doValidateAmi(
                @QueryParameter final String url,
                @QueryParameter final String accessId,
                @QueryParameter final String secretKey,
                final @QueryParameter String ami)
                throws IOException, ServletException {
            AmazonEC2 ec2 =
                OneEC2Cloud.connect(accessId, secretKey, new URL(url));
            if (ec2 != null) {
                try {
                    List<String> images = new LinkedList<String>();
                    images.add(ami);
                    List<String> owners = new LinkedList<String>();
                    List<String> users = new LinkedList<String>();
                    DescribeImagesRequest request =
                        new DescribeImagesRequest();
                    request.setImageIds(images);
                    request.setOwners(owners);
                    request.setExecutableUsers(users);
                    List<Image> img =
                        ec2.describeImages(request).getImages();
                    if (img == null || img.isEmpty()) {
                        // de-registered AMI causes an empty list to be
                        // returned. so be defensive
                        // against other possibilities
                        return FormValidation
                            .error("No such AMI, or not usable with this accessId: "
                                + ami);
                    }
                    return FormValidation.ok(img.get(0).getImageLocation()
                        + " by " + img.get(0).getImageOwnerAlias());
                } catch (AmazonClientException e) {
                    return FormValidation.error(e.getMessage());
                }
            } else {
                return FormValidation.ok(); // can't test
            }
        }

        public FormValidation doCheckIdleTerminationMinutes(
                @QueryParameter final String value) {
            if (value == null || value.trim() == "") {
                return FormValidation.ok();
            }
            try {
                int val = Integer.parseInt(value);
                if (val >= 0) {
                    return FormValidation.ok();
                }
            } catch (NumberFormatException nfe) {
            }
            return FormValidation
                .error("Idle Termination time must be a non-negative integer (or null)");
        }

        public FormValidation doCheckInstanceCapStr(
                @QueryParameter final String value) {
            if (value == null || value.trim() == "") {
                return FormValidation.ok();
            }
            try {
                int val = Integer.parseInt(value);
                if (val >= 0) {
                    return FormValidation.ok();
                }
            } catch (NumberFormatException nfe) {
            }
            return FormValidation
                .error("InstanceCap must be a non-negative integer (or null)");
        }

    }

    private void filterReservations(final List<Reservation> reservations,
            final List<Filter> filters) {

        for (Filter f : filters) {
            filterReservationsBy(reservations, f.getName(), f.getValues());
        }

    }

    private void filterReservationsBy(
            final List<Reservation> reservations, final String name,
            final List<String> values) {
        List<Reservation> emptyReservations = new ArrayList<Reservation>();

        for (Reservation reservation : reservations) {
            filterReservationBy(reservation, name, values);

            if (reservation.getInstances().size() == 0) {
                emptyReservations.add(reservation);
            }
        }

        reservations.removeAll(emptyReservations);
    }

    /**
     * @param reservation
     * @param name
     * @param values
     */
    private void filterReservationBy(final Reservation reservation,
            final String name, final List<String> values) {
        List<Instance> instances = new ArrayList<Instance>();
        for (Instance inst : reservation.getInstances()) {
            if (name.equalsIgnoreCase("image-id")) {
                String data = inst.getImageId();
                for (String value : values) {
                    if (value.equalsIgnoreCase(data)) {
                        instances.add(inst);
                        break;
                    }
                }
            }

            if (name.equalsIgnoreCase("instance-type")) {
                String data = inst.getInstanceType();
                for (String value : values) {
                    if (value.equalsIgnoreCase(data)) {
                        instances.add(inst);
                        break;
                    }
                }
            }

            if (name.equalsIgnoreCase("instance-state-name")) {
                String data = inst.getState().getName();
                for (String value : values) {
                    if (value.equalsIgnoreCase(data)) {
                        instances.add(inst);
                        break;
                    }
                }
            }

            if (name.equalsIgnoreCase("key-name")) {
                String data = inst.getKeyName();
                for (String value : values) {
                    if (value.equalsIgnoreCase(data)) {
                        instances.add(inst);
                        break;
                    }
                }
            }

            // TODO: Include more filters
            // TODO: Improve filtering mecanism

        }
        reservation.getInstances().retainAll(instances);
    }
}
