package com.workiva.om.identity.configuration.persistence;

import com.workiva.om.identity.configuration.properties.AwsWrapperProperties;
import com.workiva.om.identity.configuration.properties.HikariProperties;
import com.workiva.om.identity.configuration.properties.MySqlProperties;
import com.workiva.om.identity.infrastructure.db.ExceptionTranslationHibernateDialect;
import com.workiva.om.identity.infrastructure.db.SecurityContextAuditorAware;
import com.zaxxer.hikari.HikariDataSource;
import jakarta.persistence.EntityManagerFactory;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import javax.sql.DataSource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.flywaydb.core.Flyway;
import org.flywaydb.core.api.callback.Callback;
import org.flywaydb.core.api.migration.JavaMigration;
import org.jooq.DSLContext;
import org.jooq.SQLDialect;
import org.jooq.conf.Settings;
import org.jooq.impl.DSL;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.autoconfigure.flyway.FlywayConfigurationCustomizer;
import org.springframework.boot.autoconfigure.flyway.FlywayMigrationInitializer;
import org.springframework.boot.autoconfigure.flyway.FlywayMigrationStrategy;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.context.annotation.Profile;
import org.springframework.core.io.ResourceLoader;
import org.springframework.data.domain.AuditorAware;
import org.springframework.data.jpa.repository.config.EnableJpaAuditing;
import org.springframework.data.jpa.repository.config.EnableJpaRepositories;
import org.springframework.orm.jpa.JpaTransactionManager;
import org.springframework.orm.jpa.JpaVendorAdapter;
import org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.annotation.EnableTransactionManagement;

/**
 * Configures database components for workspaces database including
 * the auto-creation of workspaces repositories
 */
@Configuration(proxyBeanMethods = false)
@EnableTransactionManagement
@Slf4j
@EnableJpaAuditing(auditorAwareRef = "auditorProvider")
@EnableJpaRepositories(
    transactionManagerRef = "workspacesTransactionManager",
    basePackages = {
        "com.workiva.om.identity.workspaces.db.dal.model", "com.workiva.om.identity.workspaces.db.dal.repository",
    }
)
@EnableConfigurationProperties({ MySqlProperties.class, HikariProperties.class, AwsWrapperProperties.class })
public class WorkspacesDatabaseConfiguration {

    private static final Logger logger = LoggerFactory.getLogger(WorkspacesDatabaseConfiguration.class);

    public static final String WORKSPACES_FLYWAY = "flyway";
    public static final String WORKSPACES_READ_ONLY_DATASOURCE = "workspaces-read-only-dataSource";
    public static final String WORKSPACES_READ_ONLY_SLOW_DATASOURCE = "workspaces-read-only-slow-dataSource";
    public static final String WORKSPACES_PRIMARY_DATASOURCE = "workspaces-primary-dataSource";
    public static final String WORKSPACES_DATASOURCE = "workspaces-routing-dataSource";

    @Bean(name = "auditorProvider")
    public AuditorAware<String> auditorProvider() {
        return new SecurityContextAuditorAware();
    }

    @Bean
    @ConfigurationProperties("db.mysql.workspaces.primary")
    public MySqlProperties workspacesMySqlPrimaryProperties() {
        return new MySqlProperties();
    }

    @Bean
    @ConfigurationProperties("db.mysql.workspaces.read")
    public MySqlProperties workspacesMySqlReadProperties() {
        return new MySqlProperties();
    }

    @Bean
    @ConfigurationProperties("db.mysql.workspaces.read-slow")
    public MySqlProperties workspacesMySqlReadSlowProperties() {
        return new MySqlProperties();
    }

    @Bean
    @ConfigurationProperties("db.hikari.workspaces.primary")
    public HikariProperties workspacesHikariPrimaryProperties() {
        return new HikariProperties();
    }

    @Bean
    @ConfigurationProperties("db.hikari.workspaces.read")
    public HikariProperties workspacesHikariReadProperties() {
        return new HikariProperties();
    }

    @Bean
    @ConfigurationProperties("db.hikari.workspaces.read-slow")
    public HikariProperties workspacesHikariReadSlowProperties() {
        return new HikariProperties();
    }

    @Bean
    @ConfigurationProperties("db.aws.workspaces.primary")
    public AwsWrapperProperties workspacesAwsPrimaryProperties() {
        return new AwsWrapperProperties();
    }

    @Bean
    @ConfigurationProperties("db.aws.workspaces.read")
    public AwsWrapperProperties workspacesAwsReadProperties() {
        return new AwsWrapperProperties();
    }

    @Bean
    @ConfigurationProperties("db.aws.workspaces.read-slow")
    public AwsWrapperProperties workspacesAwsReadSlowProperties() {
        return new AwsWrapperProperties();
    }

    @Bean(WORKSPACES_PRIMARY_DATASOURCE)
    @Profile("!integration") // due to integration tests using h2, not aws wrapper
    public HikariDataSource workspacesPrimaryDataSource(
        @Value("${db.workspaces.url}") String url,
        @Value("${db.workspaces.username}") String username,
        @Value("${db.workspaces.password}") String password,
        @Qualifier("workspacesHikariPrimaryProperties") HikariProperties hikariProperties,
        @Qualifier("workspacesMySqlPrimaryProperties") MySqlProperties mySqlProperties,
        @Qualifier("workspacesAwsPrimaryProperties") AwsWrapperProperties awsWrapperProperties
    ) {
        return DataSourceBuilder.buildAwsWrapperDataSource(
            ConfigureAwsWrapperDataSourceCommand
                .builder()
                .jdbcUrl(url)
                .username(username)
                .password(password)
                .hikariProperties(hikariProperties)
                .mySqlProperties(mySqlProperties)
                .awsWrapperProperties(awsWrapperProperties)
                .build()
        );
    }

    @Bean(WORKSPACES_READ_ONLY_DATASOURCE)
    @Profile("!integration") // due to integration tests using h2, not aws wrapper
    public HikariDataSource workspacesReadOnlyDataSource(
        @Value("${db.workspaces.read-url}") String url,
        @Value("${db.workspaces.username}") String username,
        @Value("${db.workspaces.password}") String password,
        @Qualifier("workspacesHikariReadProperties") HikariProperties hikariProperties,
        @Qualifier("workspacesMySqlReadProperties") MySqlProperties mySqlProperties,
        @Qualifier("workspacesAwsReadProperties") AwsWrapperProperties awsWrapperProperties
    ) {
        return DataSourceBuilder.buildAwsWrapperDataSource(
            ConfigureAwsWrapperDataSourceCommand
                .builder()
                .jdbcUrl(url)
                .username(username)
                .password(password)
                .hikariProperties(hikariProperties)
                .mySqlProperties(mySqlProperties)
                .awsWrapperProperties(awsWrapperProperties)
                .build()
        );
    }

    @Bean(WORKSPACES_READ_ONLY_SLOW_DATASOURCE)
    @Profile("!integration") // due to integration tests using h2, not aws wrapper
    public HikariDataSource workspacesReadOnlySlowDataSource(
        @Value("${db.workspaces.read-url}") String url,
        @Value("${db.workspaces.username}") String username,
        @Value("${db.workspaces.password}") String password,
        @Qualifier("workspacesHikariReadSlowProperties") HikariProperties hikariProperties,
        @Qualifier("workspacesMySqlReadSlowProperties") MySqlProperties mySqlProperties,
        @Qualifier("workspacesAwsReadSlowProperties") AwsWrapperProperties awsWrapperProperties
    ) {
        return DataSourceBuilder.buildAwsWrapperDataSource(
            ConfigureAwsWrapperDataSourceCommand
                .builder()
                .jdbcUrl(url)
                .username(username)
                .password(password)
                .hikariProperties(hikariProperties)
                .mySqlProperties(mySqlProperties)
                .awsWrapperProperties(awsWrapperProperties)
                .build()
        );
    }

    @Bean(WORKSPACES_DATASOURCE)
    @Primary
    // https://fable.sh/blog/splitting-read-and-write-operations-in-spring-boot/
    public DataSource dataSource(
        @Qualifier(WORKSPACES_PRIMARY_DATASOURCE) HikariDataSource primaryDataSource,
        @Qualifier(WORKSPACES_READ_ONLY_DATASOURCE) HikariDataSource replicaDataSource,
        @Qualifier(WORKSPACES_READ_ONLY_SLOW_DATASOURCE) HikariDataSource replicaSlowDataSource
    ) {
        final RoutingDataSource routingDataSource = new RoutingDataSource();

        final Map<Object, Object> targetDataSources = new HashMap<>();
        targetDataSources.put(RoutingDataSource.Route.PRIMARY, primaryDataSource);
        targetDataSources.put(RoutingDataSource.Route.REPLICA, replicaDataSource);
        targetDataSources.put(RoutingDataSource.Route.REPLICA_SLOW, replicaSlowDataSource);

        routingDataSource.setTargetDataSources(targetDataSources);
        routingDataSource.setDefaultTargetDataSource(primaryDataSource);

        return routingDataSource;
    }

    @Bean("workspaceFlywayCustomizer")
    public FlywayConfigurationCustomizer workspaceFlywayCustomizer(
        @Value("${db.workspaces.schema}") String schema,
        @Value("${db.migration.table}") String migrationTable,
        @Value("${db.migration.locations}") String Locations
    ) {
        return configuration -> {
            Optional.ofNullable(StringUtils.trimToNull(schema)).ifPresent(configuration::schemas);
            configuration.table(migrationTable);
            configuration.locations(Locations);
            configuration.baselineOnMigrate(true);
        };
    }

    @Bean(WORKSPACES_FLYWAY)
    public Flyway flyway(
        ResourceLoader resourceLoader,
        @Qualifier(WORKSPACES_PRIMARY_DATASOURCE) ObjectProvider<DataSource> dataSource,
        @Qualifier("workspaceFlywayCustomizer") ObjectProvider<FlywayConfigurationCustomizer> workspaceFlywayCustomizer,
        ObjectProvider<JavaMigration> javaMigrations,
        ObjectProvider<Callback> callbacks
    ) {
        var flywayConfigure = Flyway
            .configure(resourceLoader.getClassLoader())
            .dataSource(dataSource.getIfAvailable())
            .javaMigrations(javaMigrations.orderedStream().toArray(JavaMigration[]::new))
            .callbacks(callbacks.orderedStream().toArray(Callback[]::new));
        workspaceFlywayCustomizer.orderedStream().forEach(customizer -> customizer.customize(flywayConfigure));
        return flywayConfigure.load();
    }

    @Bean
    public FlywayMigrationInitializer workspaceFlywayInitializer(
        @Qualifier(WORKSPACES_FLYWAY) Flyway flyway,
        @Qualifier("workspaceFlywayMigrationStrategy") ObjectProvider<FlywayMigrationStrategy> migrationStrategy
    ) {
        return new FlywayMigrationInitializer(
            flyway,
            migrationStrategy.getIfAvailable(() ->
                f -> {
                    try {
                        if ("true".equals(System.getenv("FLYWAY_REPAIR"))) {
                            f.repair();
                        }
                        f.migrate();
                    } catch (RuntimeException e) {
                        logger.error("Workspaces SQL migration failure", e);
                        throw e;
                    }
                }
            )
        );
    }

    @Bean("workspaces-create")
    public DSLContext createDSLContext(@Qualifier(WORKSPACES_PRIMARY_DATASOURCE) DataSource sqlDataSource) {
        Settings settings = new Settings().withRenderSchema(false).withRenderGroupConcatMaxLenSessionVariable(false);
        return DSL.using(sqlDataSource, SQLDialect.MYSQL, settings);
    }

    @Bean("workspaces-read")
    DSLContext readDSLContext(@Qualifier(WORKSPACES_READ_ONLY_DATASOURCE) DataSource sqlDataSource) {
        Settings settings = new Settings().withRenderSchema(false).withRenderGroupConcatMaxLenSessionVariable(false);
        return DSL.using(sqlDataSource, SQLDialect.MYSQL, settings);
    }

    @Bean
    public ExceptionTranslationHibernateDialect exceptionTranslationHibernateDialect() {
        return new ExceptionTranslationHibernateDialect();
    }

    @Bean
    @Primary
    public EntityManagerFactory entityManagerFactory(
        @Qualifier(WORKSPACES_DATASOURCE) DataSource workspacesDataSource,
        ExceptionTranslationHibernateDialect exceptionTranslationHibernateDialect,
        JpaVendorAdapter jpaVendorAdapter
    ) {
        LocalContainerEntityManagerFactoryBean emf = new LocalContainerEntityManagerFactoryBean();
        emf.setDataSource(workspacesDataSource);
        emf.setJpaVendorAdapter(jpaVendorAdapter);
        emf.setPackagesToScan(
            "com.workiva.om.identity.workspaces.db.dal.model",
            "com.workiva.administration.organization.out.rds",
            "com.workiva.administration.packages.out.rds",
            "com.workiva.administration.entity.out.rds",
            "com.workiva.administration.serviceprovider.out.rds",
            "com.workiva.administration.workspace.out.rds",
            "com.workiva.administration.workspacemembershipexpirationpolicy.out.rds",
            "com.workiva.administration.organizationmembershipexpirationpolicy.out.rds",
            "com.workiva.administration.genaiconfiguration.out.rds",
            "com.workiva.administration.stalesupportuser.out.rds",
            "com.workiva.administration.mobilesettings.out.rds",
            "com.workiva.administration.module.out.rds",
            "com.workiva.security.user.out.rds"
        );
        emf.setPersistenceUnitName("workspaces");
        emf.afterPropertiesSet();
        emf.setJpaDialect(exceptionTranslationHibernateDialect);
        return emf.getObject();
    }

    @Bean({ "transactionManager", "workspacesTransactionManager" })
    @Primary
    public PlatformTransactionManager workspacesTransactionManager(
        EntityManagerFactory entityManagerFactory,
        ExceptionTranslationHibernateDialect exceptionTranslationHibernateDialect
    ) {
        JpaTransactionManager transactionManager = new JpaTransactionManager();
        transactionManager.setEntityManagerFactory(entityManagerFactory);
        transactionManager.setJpaDialect(exceptionTranslationHibernateDialect);
        return transactionManager;
    }
}