package com.pingcap.community.blogserver.testutils;

import org.junit.jupiter.api.extension.*;
import org.springframework.context.ApplicationContext;
import org.springframework.core.io.ClassPathResource;
import org.springframework.jdbc.datasource.init.ResourceDatabasePopulator;
import org.springframework.test.context.junit.jupiter.SpringExtension;

import javax.sql.DataSource;
import java.util.Optional;

public class PopulateDataExtension implements BeforeAllCallback, AfterAllCallback {

  private DataSource dataSource;

  @Override
  public void beforeAll(ExtensionContext extensionContext) {
    ApplicationContext springContext = SpringExtension.getApplicationContext(extensionContext);
    this.dataSource = springContext.getBean(DataSource.class);
    executeScript(extensionContext, "init_data");
  }

  @Override
  public void afterAll(ExtensionContext extensionContext) {
    executeScript(extensionContext, "clear_data");
  }

  private void executeScript(ExtensionContext extensionContext, String suffix) {
    Optional<Class<?>> classOptional = extensionContext.getTestClass();
    if (classOptional.isEmpty()) return;
    Class<?> testClass = classOptional.get();
    String controllerName = camelCaseToUnderline(testClass.getSimpleName());

    ResourceDatabasePopulator populator = new ResourceDatabasePopulator();
    ClassPathResource basicDataResource = new ClassPathResource(String.format("/db/data/%s.sql", suffix));
    if (basicDataResource.exists()) {
      populator.addScript(basicDataResource);
    }
    ClassPathResource testClassDataResource = new ClassPathResource(String.format("/db/data/%s_%s.sql", controllerName, suffix));
    if (testClassDataResource.exists()) {
      populator.addScript(testClassDataResource);
    }
    populator.execute(dataSource);
  }

  private String camelCaseToUnderline(String str) {
    return str.replaceAll("([a-z])([A-Z])", "$1_$2").toLowerCase();
  }

}
