/**
 * 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.configuration.spring;

import java.util.ArrayList;
import java.util.List;

import jakarta.servlet.ServletContext;

import org.apache.ambari.server.AmbariException;
import org.apache.ambari.server.agent.AgentSessionManager;
import org.apache.ambari.server.agent.stomp.AmbariSubscriptionRegistry;
import org.apache.ambari.server.api.AmbariSendToMethodReturnValueHandler;
import org.apache.ambari.server.events.DefaultMessageEmitter;
import org.apache.ambari.server.events.publishers.AmbariEventPublisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.handler.annotation.MessageExceptionHandler;
import org.springframework.messaging.handler.invocation.HandlerMethodReturnValueHandler;
import org.springframework.messaging.simp.SimpMessagingTemplate;
import org.springframework.messaging.simp.annotation.SendToUser;
import org.springframework.messaging.simp.annotation.support.SendToMethodReturnValueHandler;
import org.springframework.messaging.simp.annotation.support.SimpAnnotationMethodMessageHandler;
import org.springframework.messaging.simp.broker.SimpleBrokerMessageHandler;
import org.springframework.messaging.support.ErrorMessage;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.socket.server.jetty.JettyRequestUpgradeStrategy;
import org.springframework.web.socket.server.support.DefaultHandshakeHandler;

import com.google.inject.Injector;

@Configuration
public class RootStompConfig {

  @Autowired
  private SimpMessagingTemplate brokerTemplate;

  private final ServletContext servletContext;

  private final org.apache.ambari.server.configuration.Configuration configuration;

  public RootStompConfig(ServletContext servletContext, Injector injector) {
    this.servletContext = servletContext;
    configuration = injector.getInstance(org.apache.ambari.server.configuration.Configuration.class);
  }

  @Bean
  public DefaultMessageEmitter defaultMessageEmitter(Injector injector) {
    org.apache.ambari.server.configuration.Configuration configuration =
        injector.getInstance(org.apache.ambari.server.configuration.Configuration.class);
    return new DefaultMessageEmitter(injector.getInstance(AgentSessionManager.class),
        brokerTemplate,
        injector.getInstance(AmbariEventPublisher.class),
        configuration.getExecutionCommandsRetryCount(),
        configuration.getExecutionCommandsRetryInterval());
  }

  @Bean
  public DefaultHandshakeHandler handshakeHandler() {
    return new DefaultHandshakeHandler(new JettyRequestUpgradeStrategy());
  }


  @Autowired
  public void configureRegistryCacheSize(SimpleBrokerMessageHandler simpleBrokerMessageHandler) throws NoSuchFieldException, IllegalAccessException {
    AmbariSubscriptionRegistry defaultSubscriptionRegistry =
        new AmbariSubscriptionRegistry(configuration.getSubscriptionRegistryCacheSize());
    simpleBrokerMessageHandler.setSubscriptionRegistry(defaultSubscriptionRegistry);
  }

  @Autowired
  public void configureGlobal(SimpAnnotationMethodMessageHandler messageHandler) {
    List<HandlerMethodReturnValueHandler> handlers = new ArrayList<>(messageHandler.getReturnValueHandlers());
    List<HandlerMethodReturnValueHandler> changedHandlers = new ArrayList<>();

    boolean handlerReplaced = false;
    for (HandlerMethodReturnValueHandler handler : handlers) {
      if (handler instanceof SendToMethodReturnValueHandler && !handlerReplaced) {
        SendToMethodReturnValueHandler sendHandler = (SendToMethodReturnValueHandler)handler;
        AmbariSendToMethodReturnValueHandler ambariSendToMethodReturnValueHandler =
          new AmbariSendToMethodReturnValueHandler(brokerTemplate, true);
        ambariSendToMethodReturnValueHandler.setHeaderInitializer(sendHandler.getHeaderInitializer());
        changedHandlers.add(ambariSendToMethodReturnValueHandler);
        handlerReplaced = true;
      } else {
        changedHandlers.add(handler);
      }
    }
    messageHandler.setReturnValueHandlers(null);
    messageHandler.setReturnValueHandlers(changedHandlers);
  }

  @ControllerAdvice
  public static class ExceptionHandlingAdvice{
    private static final Logger LOG = LoggerFactory.getLogger(ExceptionHandlingAdvice.class);

    @MessageExceptionHandler(Exception.class)
    @SendToUser("/")
    public ErrorMessage handle(Exception e) {
      LOG.error("Exception caught while processing a message: " + e.getMessage(), e);
      return new ErrorMessage(new AmbariException(e.getMessage()));
    }
  }

}
