/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.apache.ambari.server.serveraction.kerberos.stageutils;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import javax.annotation.Nullable;

import org.apache.ambari.server.AmbariException;
import org.apache.ambari.server.controller.KerberosHelper;
import org.apache.ambari.server.orm.dao.KerberosKeytabDAO;
import org.apache.ambari.server.orm.dao.KerberosKeytabPrincipalDAO;
import org.apache.ambari.server.orm.entities.KerberosKeytabEntity;
import org.apache.ambari.server.orm.entities.KerberosKeytabPrincipalEntity;
import org.apache.ambari.server.orm.entities.KerberosPrincipalEntity;
import org.apache.ambari.server.state.Cluster;
import org.apache.ambari.server.state.Clusters;
import org.apache.ambari.server.state.DesiredConfig;
import org.apache.ambari.server.state.Host;
import org.apache.ambari.server.state.Service;
import org.apache.ambari.server.state.kerberos.KerberosDescriptor;
import org.apache.ambari.server.state.kerberos.KerberosIdentityDescriptor;
import org.apache.ambari.server.utils.StageUtils;
import org.apache.commons.collections4.MapUtils;

import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import com.google.inject.Inject;
import com.google.inject.Singleton;

/**
 * Helper class to construct convenient wrappers around database entities related to kerberos.
 */
@Singleton
public class KerberosKeytabController {
  @Inject
  private KerberosKeytabDAO kerberosKeytabDAO;

  @Inject
  private KerberosKeytabPrincipalDAO kerberosKeytabPrincipalDAO;

  @Inject
  private Clusters clusters;

  //TODO: due to circular dependencies in Guice this field cannot be injected with Guice's @Inject annotation; for now we should statically inject in AmbariServer
  private static KerberosHelper kerberosHelper;

  public static void setKerberosHelper(KerberosHelper kerberosHelper) {
    KerberosKeytabController.kerberosHelper = kerberosHelper;
  }

  /**
   * Tries to find keytab by keytab path in destination filesystem.
   *
   * @param file keytab path
   * @return found keytab or null
   */
  public ResolvedKerberosKeytab getKeytabByFile(String file) {
    return getKeytabByFile(file, true);
  }

  /**
   * Tries to find keytab by keytab path in destination filesystem.
   *
   * @param file keytab path
   * @param resolvePrincipals include resolved principals
   * @return found keytab or null
   */
  public ResolvedKerberosKeytab getKeytabByFile(String file, boolean resolvePrincipals) {
    return fromKeytabEntity(kerberosKeytabDAO.find(file), resolvePrincipals, false);
  }

  /**
   * Returns all keytabs managed by ambari.
   *
   * @return all keytabs
   */
  public Set<ResolvedKerberosKeytab> getAllKeytabs() {
    return fromKeytabEntities(kerberosKeytabDAO.findAll(), false);
  }

  /**
   * Returns all keytabs that contains given principal.
   *
   * @param rkp principal to filter keytabs by
   * @return set of keytabs found
   */
  public Set<ResolvedKerberosKeytab> getFromPrincipal(ResolvedKerberosPrincipal rkp) {
    return fromKeytabEntities(kerberosKeytabDAO.findByPrincipalAndHost(rkp.getPrincipal(), rkp.getHostId()), false);
  }

  /**
   * Returns all keytabs that contains given principal without service mapping.
   *
   * @param rkp principal to filter keytabs by
   * @return set of keytabs found
   */
  public Set<ResolvedKerberosKeytab> getFromPrincipalExceptServiceMapping(ResolvedKerberosPrincipal rkp) {
    return fromKeytabEntities(kerberosKeytabDAO.findByPrincipalAndHost(rkp.getPrincipal(), rkp.getHostId()), true);
  }

  /**
   * Returns keytabs with principals filtered by host, principal name or service(and component) names.
   *
   * @param serviceComponentFilter service-component filter
   * @param hostFilter host filter
   * @param identityFilter identity(principal) filter
   * @return set of keytabs found
   */
  private Set<ResolvedKerberosKeytab> getFilteredKeytabs(Map<String, ? extends Collection<String>> serviceComponentFilter,
                                                        Set<String> hostFilter, Collection<String> identityFilter) {
    if (serviceComponentFilter == null && hostFilter == null && identityFilter == null) {
      return getAllKeytabs();
    }
    List<KerberosKeytabPrincipalDAO.KerberosKeytabPrincipalFilter> filters = splitServiceFilter(serviceComponentFilter);
    for (KerberosKeytabPrincipalDAO.KerberosKeytabPrincipalFilter filter : filters) {
      filter.setHostNames(hostFilter);
      filter.setPrincipals(identityFilter);
    }

    Set<ResolvedKerberosPrincipal> filteredPrincipals = fromPrincipalEntities(kerberosKeytabPrincipalDAO.findByFilters(filters), false);
    HashMap<String, ResolvedKerberosKeytab> resultMap = new HashMap<>();
    for (ResolvedKerberosPrincipal principal : filteredPrincipals) {
      if (!resultMap.containsKey(principal.getKeytabPath())) {
        resultMap.put(principal.getKeytabPath(), getKeytabByFile(principal.getKeytabPath(), false));
      }
      ResolvedKerberosKeytab keytab = resultMap.get(principal.getKeytabPath());
      keytab.addPrincipal(principal);
    }
    return Sets.newHashSet(resultMap.values());
  }

  public Set<ResolvedKerberosKeytab> getFilteredKeytabs(Collection<KerberosIdentityDescriptor> serviceIdentities, Set<String> hostFilter, Collection<String> identityFilters) {
    final Collection<String> enhancedIdentityFilters = populateIdentityFilter(identityFilters, serviceIdentities);
    return getFilteredKeytabs((Map<String, ? extends Collection<String>>) null, hostFilter, enhancedIdentityFilters);
  }

  /**
   * This function split serviceComponentFilter to two filters, one with specific components, and another one with service
   * only. Can return only one filter if filter contain only one type of mapping(whole service or component based)
   * or empty filter if no serviceComponentFilter provided.
   *
   * @param serviceComponentFilter
   * @return
   */
  private List<KerberosKeytabPrincipalDAO.KerberosKeytabPrincipalFilter> splitServiceFilter(Map<String, ? extends Collection<String>> serviceComponentFilter) {
    if (serviceComponentFilter != null && serviceComponentFilter.size() > 0) {
      Set<String> serviceSet = new HashSet<>();
      Set<String> componentSet = new HashSet<>();
      Set<String> serviceOnlySet = new HashSet<>();

      // Split the filter into a service/component filter or a service-only filter.
      for (Map.Entry<String, ? extends Collection<String>> entry : serviceComponentFilter.entrySet()) {
        String serviceName = entry.getKey();
        Collection<String> serviceComponents = entry.getValue();

        if((serviceComponents == null) || serviceComponents.contains("*")) {
          serviceOnlySet.add(serviceName);
        }
        else {
          serviceSet.add(serviceName);
          componentSet.addAll(serviceComponents);
        }
      }

      List<KerberosKeytabPrincipalDAO.KerberosKeytabPrincipalFilter> result = new ArrayList<>();
      // Handle the service/component filter
      if (serviceSet.size() > 0) {
        result.add(KerberosKeytabPrincipalDAO.KerberosKeytabPrincipalFilter.createFilter(
          null,
          serviceSet,
          componentSet,
          null
        ));
      }
      // Handler the service/* filter
      if (serviceOnlySet.size() > 0) {
        result.add(KerberosKeytabPrincipalDAO.KerberosKeytabPrincipalFilter.createFilter(
          null,
          serviceOnlySet,
          null,
          null
        ));
      }
      if (result.size() > 0) {
        return result;
      }
    }

    return Lists.newArrayList(KerberosKeytabPrincipalDAO.KerberosKeytabPrincipalFilter.createEmptyFilter());
  }

  private ResolvedKerberosKeytab fromKeytabEntity(KerberosKeytabEntity kke, boolean resolvePrincipals, boolean exceptServiceMapping) {
    Set<ResolvedKerberosPrincipal> principals = resolvePrincipals ?
        fromPrincipalEntities(kke.getKerberosKeytabPrincipalEntities(), exceptServiceMapping) : new HashSet<>();
    return new ResolvedKerberosKeytab(
      kke.getKeytabPath(),
      kke.getOwnerName(),
      kke.getOwnerAccess(),
      kke.getGroupName(),
      kke.getGroupAccess(),
      principals,
      kke.isAmbariServerKeytab(),
      kke.isWriteAmbariJaasFile()
    );
  }

  private ResolvedKerberosKeytab fromKeytabEntity(KerberosKeytabEntity kke) {
    return fromKeytabEntity(kke, true, false);
  }

  private Set<ResolvedKerberosKeytab> fromKeytabEntities(Collection<KerberosKeytabEntity> keytabEntities, boolean exceptServiceMapping) {
    ImmutableSet.Builder<ResolvedKerberosKeytab> builder = ImmutableSet.builder();
    for (KerberosKeytabEntity kke : keytabEntities) {
      builder.add(fromKeytabEntity(kke, true, exceptServiceMapping));
    }
    return builder.build();
  }

  private Set<ResolvedKerberosPrincipal> fromPrincipalEntities(Collection<KerberosKeytabPrincipalEntity> principalEntities, boolean exceptServiceMapping) {
    ImmutableSet.Builder<ResolvedKerberosPrincipal> builder = ImmutableSet.builder();
    for (KerberosKeytabPrincipalEntity kkpe : principalEntities) {
      KerberosPrincipalEntity kpe = kkpe.getKerberosPrincipalEntity();

      if(kpe != null) {
        ResolvedKerberosPrincipal rkp = new ResolvedKerberosPrincipal(
            kkpe.getHostId(),
            kkpe.getHostName(),
            kkpe.getPrincipalName(),
            kpe.isService(),
            kpe.getCachedKeytabPath(),
            kkpe.getKeytabPath(),
            exceptServiceMapping ? null : kkpe.getServiceMappingAsMultimap());
        builder.add(rkp);
      }
    }
    return builder.build();
  }

  /**
   * Adjust service component filter according to installed services
   *
   * @param cluster                cluster
   * @param includeAmbariAsService
   * @param serviceComponentFilter
   * @return
   * @throws AmbariException
   */
  public Map<String, Collection<String>> adjustServiceComponentFilter(Cluster cluster, boolean includeAmbariAsService, Map<String, ? extends Collection<String>> serviceComponentFilter) throws AmbariException {
    Map<String, Collection<String>> adjustedFilter = new HashMap<>();

    Map<String, Service> installedServices = (cluster == null) ? null : cluster.getServices();
    if(includeAmbariAsService) {
      installedServices = (installedServices == null) ? new HashMap<>() : new HashMap<>(installedServices);
      installedServices.put("AMBARI", null);
    }

    if (!MapUtils.isEmpty(installedServices)) {
      if (serviceComponentFilter != null) {
        // prune off services that are not installed, or considered installed - like AMBARI
        for (Map.Entry<String, ? extends Collection<String>> entry : serviceComponentFilter.entrySet()) {
          String serviceName = entry.getKey();

          if (installedServices.containsKey(serviceName)) {
            adjustedFilter.put(serviceName, entry.getValue());
          }
        }
      } else {
        // return only the set of installed services
        for (String serviceName : installedServices.keySet()) {
          // Add an entry to indicate the service and all of it's components should be considered
          adjustedFilter.put(serviceName, Collections.singletonList("*"));
        }
      }
    }

    return adjustedFilter;
  }

  public Collection<KerberosIdentityDescriptor> getServiceIdentities(String clusterName, Collection<String> services,
                                                                     @Nullable Map<String, DesiredConfig> desiredConfigs)
    throws AmbariException {

    final Collection<KerberosIdentityDescriptor> serviceIdentities = new ArrayList<>();
    Cluster cluster = clusters.getCluster(clusterName);

    if (desiredConfigs == null) {
      desiredConfigs = cluster.getDesiredConfigs();
    }

    KerberosDescriptor kerberosDescriptor = kerberosHelper.getKerberosDescriptor(cluster, false);
    KerberosDescriptor userDescriptor = kerberosHelper.getKerberosDescriptorUpdates(cluster);
    Map<String, Map<String, Map<String, String>>> hostConfigurations = new HashMap<>();
    Map<String, Host> hostMap = clusters.getHostsForCluster(clusterName);
    Set<String> hosts = new HashSet<>(hostMap.keySet());
    Map<String, String> componentHosts = new HashMap<>();

    for (String hostName: hosts) {
      hostConfigurations.put(
        hostName,
        kerberosHelper.calculateConfigurations(cluster, hostName, kerberosDescriptor, userDescriptor,
          false,false, componentHosts, desiredConfigs)
      );
    }

    // ambari-agent may not be installed on ambari-server. Add ambari-server after calculating configurations.
    // If not, HostNotFoundException will raise
    String ambariServerHostname = StageUtils.getHostName();
    hosts.add(ambariServerHostname);

    for (String service: services) {
      kerberosHelper.getActiveIdentities(clusterName,null, service, null,true,
        hostConfigurations,kerberosDescriptor, desiredConfigs).values().forEach(serviceIdentities::addAll);
    }
    return serviceIdentities;
  }

  private Collection<String> populateIdentityFilter(Collection<String> identityFilters, Collection<KerberosIdentityDescriptor> serviceIdentities) {
    if (serviceIdentities != null) {
      identityFilters = identityFilters == null ? new HashSet<>() : identityFilters;
      for (KerberosIdentityDescriptor serviceIdentity : serviceIdentities) {
        if (!KerberosHelper.AMBARI_SERVER_KERBEROS_IDENTITY_NAME.equals(serviceIdentity.getName())) {
          identityFilters.add(serviceIdentity.getPrincipalDescriptor().getName());
        }
      }
    }
    return identityFilters;
  }
}
