/*
 * Copyright The OpenTelemetry Authors
 * SPDX-License-Identifier: Apache-2.0
 */

package io.opentelemetry.javaagent.instrumentation.spring.rmi.v4_0;

import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.equalTo;
import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.satisfies;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;

import io.opentelemetry.api.trace.SpanKind;
import io.opentelemetry.instrumentation.test.utils.PortUtils;
import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
import io.opentelemetry.sdk.testing.assertj.SpanDataAssert;
import io.opentelemetry.sdk.trace.data.StatusData;
import io.opentelemetry.semconv.ExceptionAttributes;
import io.opentelemetry.semconv.incubating.RpcIncubatingAttributes;
import java.rmi.RemoteException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
import javax.ejb.EJBException;
import javax.ejb.embeddable.EJBContainer;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.extension.RegisterExtension;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.EnumSource;
import org.springframework.boot.SpringApplication;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ConfigurableApplicationContext;
import org.springframework.context.annotation.Bean;
import org.springframework.context.support.ClassPathXmlApplicationContext;
import org.springframework.remoting.rmi.RmiProxyFactoryBean;
import org.springframework.remoting.rmi.RmiServiceExporter;
import org.springframework.remoting.support.RemoteExporter;
import org.springframework.stereotype.Component;
import springrmi.app.SpringRmiGreeter;
import springrmi.app.SpringRmiGreeterImpl;

class SpringRmiTest {

  @RegisterExtension
  static final InstrumentationExtension testing = AgentInstrumentationExtension.create();

  private static ConfigurableApplicationContext serverAppContext;
  private static ConfigurableApplicationContext clientAppContext;
  private static ConfigurableApplicationContext xmlAppContext;
  private static EJBContainer ejbContainer;

  static int registryPort;

  @Component
  private static class ServerConfig {
    @Bean
    static RemoteExporter registerRmiExporter() {
      SpringRmiGreeter greeter = new SpringRmiGreeterImpl();

      RmiServiceExporter exporter = new RmiServiceExporter();
      exporter.setServiceName("springRmiGreeter");
      exporter.setServiceInterface(SpringRmiGreeter.class);
      exporter.setService(greeter);
      exporter.setRegistryPort(registryPort);

      return exporter;
    }
  }

  @Component
  private static class ClientConfig {
    @Bean
    static RmiProxyFactoryBean rmiProxy() {
      RmiProxyFactoryBean bean = new RmiProxyFactoryBean();
      bean.setServiceInterface(SpringRmiGreeter.class);
      bean.setServiceUrl("rmi://localhost:" + registryPort + "/springRmiGreeter");
      return bean;
    }
  }

  @BeforeAll
  static void beforeAll() throws Exception {
    registryPort = PortUtils.findOpenPort();

    Map<String, Object> map = new HashMap<>();
    map.put(EJBContainer.APP_NAME, "test");
    map.put(EJBContainer.MODULES, new java.io.File("build/classes/java/test"));
    ejbContainer = EJBContainer.createEJBContainer(map);

    Map<String, Object> props = new HashMap<>();
    props.put("spring.jmx.enabled", false);
    props.put("spring.main.web-application-type", "none");

    SpringApplication serverApp = new SpringApplication(ServerConfig.class);
    serverApp.setDefaultProperties(props);
    serverAppContext = serverApp.run();

    SpringApplication clientApp = new SpringApplication(ClientConfig.class);
    clientApp.setDefaultProperties(props);
    clientAppContext = clientApp.run();

    xmlAppContext = new ClassPathXmlApplicationContext("spring-rmi.xml");
  }

  @AfterAll
  static void afterAll() {
    serverAppContext.close();
    clientAppContext.close();
    xmlAppContext.close();
    ejbContainer.close();
  }

  @SuppressWarnings("ImmutableEnumChecker")
  private enum TestSource {
    RMI(
        clientAppContext,
        "springrmi.app.SpringRmiGreeterImpl",
        "spring_rmi",
        IllegalStateException.class),
    EJB(xmlAppContext, "springrmi.app.ejb.SpringRmiGreeterRemote", "java_rmi", EJBException.class);

    final ApplicationContext appContext;
    final String remoteClassName;
    final String serverSystem;
    final Class<? extends Throwable> expectedException;

    TestSource(
        ApplicationContext appContext,
        String remoteClassName,
        String serverSystem,
        Class<? extends Throwable> expectedException) {
      this.appContext = appContext;
      this.remoteClassName = remoteClassName;
      this.serverSystem = serverSystem;
      this.expectedException = expectedException;
    }
  }

  @ParameterizedTest(autoCloseArguments = false)
  @EnumSource(TestSource.class)
  void clientCallCreatesSpans(TestSource testSource) throws RemoteException {
    SpringRmiGreeter client = testSource.appContext.getBean(SpringRmiGreeter.class);
    String response = testing.runWithSpan("parent", () -> client.hello("Test Name"));
    assertEquals(response, "Hello Test Name");
    testing.waitAndAssertTraces(
        trace -> {
          List<Consumer<SpanDataAssert>> assertions = new ArrayList<>();
          assertions.add(span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent());
          assertions.add(
              span ->
                  span.hasName("springrmi.app.SpringRmiGreeter/hello")
                      .hasKind(SpanKind.CLIENT)
                      .hasParent(trace.getSpan(0))
                      .hasAttributesSatisfying(
                          equalTo(RpcIncubatingAttributes.RPC_SYSTEM, "spring_rmi"),
                          equalTo(
                              RpcIncubatingAttributes.RPC_SERVICE,
                              "springrmi.app.SpringRmiGreeter"),
                          equalTo(RpcIncubatingAttributes.RPC_METHOD, "hello")));
          if (testSource == TestSource.RMI) {
            assertions.add(
                span ->
                    span.hasName(testSource.remoteClassName + "/hello")
                        .hasKind(SpanKind.SERVER)
                        .hasParent(trace.getSpan(1))
                        .hasAttributesSatisfying(
                            equalTo(RpcIncubatingAttributes.RPC_SYSTEM, testSource.serverSystem),
                            equalTo(
                                RpcIncubatingAttributes.RPC_SERVICE, testSource.remoteClassName),
                            equalTo(RpcIncubatingAttributes.RPC_METHOD, "hello")));
          }

          trace.hasSpansSatisfyingExactly(assertions);
        });
  }

  @ParameterizedTest(autoCloseArguments = false)
  @EnumSource(TestSource.class)
  void throwsException(TestSource testSource) {
    SpringRmiGreeter client = testSource.appContext.getBean(SpringRmiGreeter.class);
    Throwable error =
        assertThrows(
            testSource.expectedException, () -> testing.runWithSpan("parent", client::exceptional));
    testing.waitAndAssertTraces(
        trace -> {
          List<Consumer<SpanDataAssert>> assertions = new ArrayList<>();
          assertions.add(
              span ->
                  span.hasName("parent")
                      .hasKind(SpanKind.INTERNAL)
                      .hasStatus(StatusData.error())
                      .hasNoParent()
                      .hasEventsSatisfyingExactly(
                          event ->
                              event
                                  .hasName("exception")
                                  .hasAttributesSatisfying(
                                      equalTo(
                                          ExceptionAttributes.EXCEPTION_TYPE,
                                          error.getClass().getCanonicalName()),
                                      equalTo(
                                          ExceptionAttributes.EXCEPTION_MESSAGE,
                                          error.getMessage()),
                                      satisfies(
                                          ExceptionAttributes.EXCEPTION_STACKTRACE,
                                          val -> val.isInstanceOf(String.class)))));
          assertions.add(
              span ->
                  span.hasName("springrmi.app.SpringRmiGreeter/exceptional")
                      .hasKind(SpanKind.CLIENT)
                      .hasStatus(StatusData.error())
                      .hasParent(trace.getSpan(0))
                      .hasEventsSatisfyingExactly(
                          event ->
                              event
                                  .hasName("exception")
                                  .hasAttributesSatisfying(
                                      equalTo(
                                          ExceptionAttributes.EXCEPTION_TYPE,
                                          error.getClass().getCanonicalName()),
                                      equalTo(
                                          ExceptionAttributes.EXCEPTION_MESSAGE,
                                          error.getMessage()),
                                      satisfies(
                                          ExceptionAttributes.EXCEPTION_STACKTRACE,
                                          val -> val.isInstanceOf(String.class))))
                      .hasAttributesSatisfying(
                          equalTo(RpcIncubatingAttributes.RPC_SYSTEM, "spring_rmi"),
                          equalTo(
                              RpcIncubatingAttributes.RPC_SERVICE,
                              "springrmi.app.SpringRmiGreeter"),
                          equalTo(RpcIncubatingAttributes.RPC_METHOD, "exceptional")));
          if (testSource == TestSource.RMI) {
            assertions.add(
                span ->
                    span.hasName(testSource.remoteClassName + "/exceptional")
                        .hasKind(SpanKind.SERVER)
                        .hasParent(trace.getSpan(1))
                        .hasStatus(StatusData.error())
                        .hasEventsSatisfyingExactly(
                            event ->
                                event
                                    .hasName("exception")
                                    .hasAttributesSatisfying(
                                        equalTo(
                                            ExceptionAttributes.EXCEPTION_TYPE,
                                            error.getClass().getCanonicalName()),
                                        equalTo(
                                            ExceptionAttributes.EXCEPTION_MESSAGE,
                                            error.getMessage()),
                                        satisfies(
                                            ExceptionAttributes.EXCEPTION_STACKTRACE,
                                            val -> val.isInstanceOf(String.class))))
                        .hasAttributesSatisfying(
                            equalTo(RpcIncubatingAttributes.RPC_SYSTEM, testSource.serverSystem),
                            equalTo(
                                RpcIncubatingAttributes.RPC_SERVICE, testSource.remoteClassName),
                            equalTo(RpcIncubatingAttributes.RPC_METHOD, "exceptional")));
          }

          trace.hasSpansSatisfyingExactly(assertions);
        });
  }
}
