/**
 * 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.agent.stomp;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

import jakarta.ws.rs.WebApplicationException;

import org.apache.ambari.server.AmbariException;
import org.apache.ambari.server.agent.AgentSessionManager;
import org.apache.ambari.server.agent.HeartBeat;
import org.apache.ambari.server.agent.HeartBeatHandler;
import org.apache.ambari.server.agent.HeartBeatResponse;
import org.apache.ambari.server.agent.Register;
import org.apache.ambari.server.agent.RegistrationResponse;
import org.apache.ambari.server.agent.RegistrationStatus;
import org.apache.ambari.server.configuration.Configuration;
import org.apache.ambari.server.configuration.spring.GuiceBeansConfig;
import org.apache.ambari.server.state.cluster.ClustersImpl;
import org.apache.ambari.server.state.fsm.InvalidStateTransitionException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Import;
import org.springframework.messaging.handler.annotation.Header;
import org.springframework.messaging.handler.annotation.MessageMapping;
import org.springframework.messaging.simp.annotation.SendToUser;
import org.springframework.stereotype.Controller;

import com.google.common.util.concurrent.ThreadFactoryBuilder;
import com.google.inject.Injector;
import com.google.inject.persist.UnitOfWork;

@Controller
@SendToUser("/")
@MessageMapping("/")
@Import(GuiceBeansConfig.class)
public class HeartbeatController {
  private static Logger LOG = LoggerFactory.getLogger(HeartbeatController.class);
  private final HeartBeatHandler hh;
  private final ClustersImpl clusters;
  private final AgentSessionManager agentSessionManager;
  private final LinkedBlockingQueue queue;
  private final ThreadFactory threadFactoryExecutor = new ThreadFactoryBuilder().setNameFormat("agent-register-processor-%d").build();
  private final ThreadFactory threadFactoryTimeout = new ThreadFactoryBuilder().setNameFormat("agent-register-timeout-%d").build();
  private final ExecutorService executor;
  private final ScheduledExecutorService scheduledExecutorService;
  private final UnitOfWork unitOfWork;

  @Autowired
  private AgentsRegistrationQueue agentsRegistrationQueue;

  public HeartbeatController(Injector injector) {
    hh = injector.getInstance(HeartBeatHandler.class);
    clusters = injector.getInstance(ClustersImpl.class);
    unitOfWork = injector.getInstance(UnitOfWork.class);
    agentSessionManager = injector.getInstance(AgentSessionManager.class);

    Configuration configuration = injector.getInstance(Configuration.class);
    queue = new LinkedBlockingQueue(configuration.getAgentsRegistrationQueueSize());
    executor = new ThreadPoolExecutor(configuration.getRegistrationThreadPoolSize(),
        configuration.getRegistrationThreadPoolSize(), 0L, TimeUnit.MILLISECONDS, queue, threadFactoryExecutor);
    scheduledExecutorService = Executors.newScheduledThreadPool(1, threadFactoryTimeout);
  }

  @MessageMapping("/register")
  public CompletableFuture<RegistrationResponse> register(@Header String simpSessionId, Register message)
      throws WebApplicationException, InvalidStateTransitionException, AmbariException {
    CompletableFuture<RegistrationResponse> completableFuture = new CompletableFuture<>();

    Future<RegistrationResponse> future = executor.submit(() -> {
      try {
        unitOfWork.begin();
        RegistrationResponse response = null;
        try {
          /* Call into the heartbeat handler */
          response = hh.handleRegistration(message);
          agentSessionManager.register(simpSessionId,
              clusters.getHost(message.getHostname()));
          LOG.debug("Sending registration response " + response);
        } catch (Exception ex) {
          LOG.info(ex.getMessage(), ex);
          response = new RegistrationResponse();
          response.setResponseId(-1);
          response.setResponseStatus(RegistrationStatus.FAILED);
          response.setExitstatus(1);
          response.setLog(ex.getMessage());
          completableFuture.complete(response);
          return response;
        }
        completableFuture.complete(response);
        return response;
      } finally {
        unitOfWork.end();
      }
    });

    scheduledExecutorService.schedule(new RegistrationTimeoutTask(future, completableFuture), 8, TimeUnit.SECONDS);
    return completableFuture;
  }

  @MessageMapping("/heartbeat")
  public HeartBeatResponse heartbeat(@Header String simpSessionId, HeartBeat message) {
    try {
      unitOfWork.begin();
      if (LOG.isDebugEnabled()) {
        LOG.debug("Received Heartbeat message " + message);
      }
      HeartBeatResponse heartBeatResponse;
      try {
        if (!agentSessionManager.isRegistered(simpSessionId)) {
          //Server restarted, or unknown host.
          LOG.error(String.format("Host with [%s] sessionId not registered", simpSessionId));
          return hh.createRegisterCommand();
        }
        message.setHostname(agentSessionManager.getHost(simpSessionId).getHostName());
        heartBeatResponse = hh.handleHeartBeat(message);
        agentsRegistrationQueue.complete(simpSessionId);
        if (LOG.isDebugEnabled()) {
          LOG.debug("Sending heartbeat response with response id " + heartBeatResponse.getResponseId());
          LOG.debug("Response details " + heartBeatResponse);
        }
      } catch (Exception e) {
        LOG.warn("Error in HeartBeat", e);
        throw new WebApplicationException(500);
      }
      return heartBeatResponse;
    } finally {
      unitOfWork.end();
    }
  }

  private class RegistrationTimeoutTask implements Runnable {
    private Future<RegistrationResponse> task;
    private CompletableFuture<RegistrationResponse> completableFuture;

    public RegistrationTimeoutTask(Future<RegistrationResponse> task, CompletableFuture<RegistrationResponse> completableFuture) {

      this.task = task;
      this.completableFuture = completableFuture;
    }

    @Override
    public void run() {
      boolean cancelled = task.cancel(false);
      if (cancelled) {
        completableFuture.cancel(false);
      }
    }
  }
}
