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

package io.opentelemetry.instrumentation.r2dbc.v1_0;

import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.equalTo;
import static io.opentelemetry.semconv.trace.attributes.SemanticAttributes.DB_CONNECTION_STRING;
import static io.opentelemetry.semconv.trace.attributes.SemanticAttributes.DB_NAME;
import static io.opentelemetry.semconv.trace.attributes.SemanticAttributes.DB_OPERATION;
import static io.opentelemetry.semconv.trace.attributes.SemanticAttributes.DB_SQL_TABLE;
import static io.opentelemetry.semconv.trace.attributes.SemanticAttributes.DB_STATEMENT;
import static io.opentelemetry.semconv.trace.attributes.SemanticAttributes.DB_SYSTEM;
import static io.opentelemetry.semconv.trace.attributes.SemanticAttributes.DB_USER;
import static io.opentelemetry.semconv.trace.attributes.SemanticAttributes.NET_PEER_NAME;
import static io.opentelemetry.semconv.trace.attributes.SemanticAttributes.NET_PEER_PORT;
import static io.r2dbc.spi.ConnectionFactoryOptions.DATABASE;
import static io.r2dbc.spi.ConnectionFactoryOptions.DRIVER;
import static io.r2dbc.spi.ConnectionFactoryOptions.HOST;
import static io.r2dbc.spi.ConnectionFactoryOptions.PASSWORD;
import static io.r2dbc.spi.ConnectionFactoryOptions.PORT;
import static io.r2dbc.spi.ConnectionFactoryOptions.USER;
import static org.junit.jupiter.api.Named.named;

import com.google.errorprone.annotations.CanIgnoreReturnValue;
import io.opentelemetry.api.trace.SpanKind;
import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
import io.r2dbc.spi.ConnectionFactories;
import io.r2dbc.spi.ConnectionFactory;
import io.r2dbc.spi.ConnectionFactoryOptions;
import java.time.Duration;
import java.util.HashMap;
import java.util.Map;
import java.util.stream.Stream;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.TestInstance;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.containers.output.Slf4jLogConsumer;
import org.testcontainers.containers.wait.strategy.Wait;
import reactor.core.publisher.Mono;

@TestInstance(TestInstance.Lifecycle.PER_CLASS)
public abstract class AbstractR2dbcStatementTest {
  private static final Logger logger = LoggerFactory.getLogger(AbstractR2dbcStatementTest.class);

  protected abstract InstrumentationExtension getTesting();

  private static final String USER_DB = "SA";
  private static final String PW_DB = "password123";
  private static final String DB = "tempdb";

  private static final DbSystemProps POSTGRESQL =
      new DbSystemProps("postgresql", "postgres:9.6.8", 5432)
          .envVariables(
              "POSTGRES_USER", USER_DB,
              "POSTGRES_PASSWORD", PW_DB,
              "POSTGRES_DB", DB);

  private static final DbSystemProps MARIADB =
      new DbSystemProps("mariadb", "mariadb:10.3.6", 3306)
          .envVariables(
              "MYSQL_ROOT_PASSWORD", PW_DB,
              "MYSQL_USER", USER_DB,
              "MYSQL_PASSWORD", PW_DB,
              "MYSQL_DATABASE", DB);

  private static final DbSystemProps MYSQL =
      new DbSystemProps("mysql", "mysql:8.0.32", 3306)
          .envVariables(
              "MYSQL_ROOT_PASSWORD", PW_DB,
              "MYSQL_USER", USER_DB,
              "MYSQL_PASSWORD", PW_DB,
              "MYSQL_DATABASE", DB);

  private static final Map<String, DbSystemProps> SYSTEMS = new HashMap<>();

  static {
    SYSTEMS.put(POSTGRESQL.system, POSTGRESQL);
    SYSTEMS.put(MYSQL.system, MYSQL);
    SYSTEMS.put(MARIADB.system, MARIADB);
  }

  private static Integer port;
  private static GenericContainer<?> container;

  protected ConnectionFactory createProxyConnectionFactory(
      ConnectionFactoryOptions connectionFactoryOptions) {
    return ConnectionFactories.find(connectionFactoryOptions);
  }

  @AfterAll
  void stopContainer() {
    if (container != null) {
      container.stop();
    }
  }

  void startContainer(DbSystemProps props) {
    if (container != null && container.getDockerImageName().equals(props.image)) {
      return;
    }
    if (container != null) {
      container.stop();
    }
    if (props.image != null) {
      container =
          new GenericContainer<>(props.image)
              .withEnv(props.envVariables)
              .withExposedPorts(props.port)
              .withLogConsumer(new Slf4jLogConsumer(logger))
              .withStartupTimeout(Duration.ofMinutes(2));
      if (props == POSTGRESQL) {
        container.waitingFor(
            Wait.forLogMessage(".*database system is ready to accept connections.*", 2));
      }
      container.start();
      port = container.getMappedPort(props.port);
    }
  }

  @ParameterizedTest(name = "{index}: {0}")
  @MethodSource("provideParameters")
  void testQueries(Parameter parameter) {
    DbSystemProps props = SYSTEMS.get(parameter.system);
    startContainer(props);
    ConnectionFactory connectionFactory =
        createProxyConnectionFactory(
            ConnectionFactoryOptions.builder()
                .option(DRIVER, props.system)
                .option(HOST, "localhost")
                .option(PORT, port)
                .option(USER, USER_DB)
                .option(PASSWORD, PW_DB)
                .option(DATABASE, DB)
                .build());

    getTesting()
        .runWithSpan(
            "parent",
            () -> {
              Mono.from(connectionFactory.create())
                  .flatMapMany(
                      connection ->
                          Mono.from(connection.createStatement(parameter.statement).execute())
                              // Subscribe to the Statement.execute()
                              .flatMapMany(result -> result.map((row, metadata) -> ""))
                              .concatWith(Mono.from(connection.close()).cast(String.class)))
                  .doFinally(e -> getTesting().runWithSpan("child", () -> {}))
                  .blockLast(Duration.ofMinutes(1));
            });

    getTesting()
        .waitAndAssertTraces(
            trace ->
                trace.hasSpansSatisfyingExactly(
                    span -> span.hasName("parent").hasKind(SpanKind.INTERNAL),
                    span ->
                        span.hasName(parameter.spanName)
                            .hasKind(SpanKind.CLIENT)
                            .hasParent(trace.getSpan(0))
                            .hasAttributesSatisfyingExactly(
                                equalTo(
                                    DB_CONNECTION_STRING,
                                    parameter.system + "://localhost:" + port),
                                equalTo(DB_SYSTEM, parameter.system),
                                equalTo(DB_NAME, DB),
                                equalTo(DB_USER, USER_DB),
                                equalTo(DB_STATEMENT, parameter.expectedStatement),
                                equalTo(DB_OPERATION, parameter.operation),
                                equalTo(DB_SQL_TABLE, parameter.table),
                                equalTo(NET_PEER_NAME, "localhost"),
                                equalTo(NET_PEER_PORT, port)),
                    span ->
                        span.hasName("child")
                            .hasKind(SpanKind.INTERNAL)
                            .hasParent(trace.getSpan(0))));
  }

  private static Stream<Arguments> provideParameters() {
    return SYSTEMS.values().stream()
        .flatMap(
            system ->
                Stream.of(
                    Arguments.of(
                        named(
                            system.system + " Simple Select",
                            new Parameter(
                                system.system,
                                "SELECT 3",
                                "SELECT ?",
                                "SELECT " + DB,
                                null,
                                "SELECT"))),
                    Arguments.of(
                        named(
                            system.system + " Create Table",
                            new Parameter(
                                system.system,
                                "CREATE TABLE person (id SERIAL PRIMARY KEY, first_name VARCHAR(255), last_name VARCHAR(255))",
                                "CREATE TABLE person (id SERIAL PRIMARY KEY, first_name VARCHAR(?), last_name VARCHAR(?))",
                                DB,
                                null,
                                null))),
                    Arguments.of(
                        named(
                            system.system + " Insert",
                            new Parameter(
                                system.system,
                                "INSERT INTO person (id, first_name, last_name) values (1, 'tom', 'johnson')",
                                "INSERT INTO person (id, first_name, last_name) values (?, ?, ?)",
                                "INSERT " + DB + ".person",
                                "person",
                                "INSERT"))),
                    Arguments.of(
                        named(
                            system.system + " Select from Table",
                            new Parameter(
                                system.system,
                                "SELECT * FROM person where first_name = 'tom'",
                                "SELECT * FROM person where first_name = ?",
                                "SELECT " + DB + ".person",
                                "person",
                                "SELECT")))));
  }

  private static class Parameter {

    public final String system;
    public final String statement;
    public final String expectedStatement;
    public final String spanName;
    public final String table;
    public final String operation;

    public Parameter(
        String system,
        String statement,
        String expectedStatement,
        String spanName,
        String table,
        String operation) {
      this.system = system;
      this.statement = statement;
      this.expectedStatement = expectedStatement;
      this.spanName = spanName;
      this.table = table;
      this.operation = operation;
    }
  }

  private static class DbSystemProps {
    public final String system;
    public final String image;
    public final int port;
    public final Map<String, String> envVariables = new HashMap<>();

    public DbSystemProps(String system, String image, int port) {
      this.system = system;
      this.image = image;
      this.port = port;
    }

    @CanIgnoreReturnValue
    public DbSystemProps envVariables(String... keyValues) {
      for (int i = 0; i < keyValues.length / 2; i++) {
        envVariables.put(keyValues[2 * i], keyValues[2 * i + 1]);
      }
      return this;
    }
  }
}
