// Copyright (c) Philipp Wagner. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

package de.bytefish.multitenancy.routing;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.zaxxer.hikari.HikariDataSource;
import de.bytefish.multitenancy.core.ThreadLocalStorage;
import de.bytefish.multitenancy.routing.config.DatabaseConfiguration;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.scheduling.annotation.Scheduled;

import javax.sql.DataSource;

import java.io.File;
import java.util.*;
import java.util.concurrent.ConcurrentMap;
import java.util.stream.Collectors;

public class DynamicTenantAwareRoutingSource extends AbstractRoutingDataSource {

    private final String filename;
    private final ObjectMapper objectMapper;
    private final ConcurrentMap<String, HikariDataSource> tenants;

    public DynamicTenantAwareRoutingSource(String filename) {
        this(filename, new ObjectMapper());
    }

    public DynamicTenantAwareRoutingSource(String filename, ObjectMapper objectMapper) {
        this.filename = filename;
        this.objectMapper = objectMapper;
        this.tenants = getDataSources();
    }

    @Override
    public void afterPropertiesSet() {
        // Nothing to do ..
    }

    @Override
    protected DataSource determineTargetDataSource() {
        String lookupKey = (String) determineCurrentLookupKey();

        // And finally return it:
        return tenants.get(lookupKey);
    }

    @Override
    protected Object determineCurrentLookupKey() {
        return ThreadLocalStorage.getTenantName();
    }

    private ConcurrentMap<String, HikariDataSource> getDataSources() {

        // Deserialize the JSON:
        DatabaseConfiguration[] configurations = getDatabaseConfigurations();

        // Now create a Lookup Table:
        return Arrays
                .stream(configurations)
                .collect(Collectors.toConcurrentMap(x -> x.getTenant(), x -> buildDataSource(x)));
    }

    private DatabaseConfiguration[] getDatabaseConfigurations() {
        try {
            return objectMapper.readValue(new File(filename), DatabaseConfiguration[].class);
        } catch(Exception e) {
            throw new RuntimeException(e);
        }
    }

    private HikariDataSource buildDataSource(DatabaseConfiguration configuration) {
        HikariDataSource dataSource = new HikariDataSource();

        dataSource.setInitializationFailTimeout(0);
        dataSource.setMaximumPoolSize(5);
        dataSource.setDataSourceClassName(configuration.getDataSourceClassName());
        dataSource.addDataSourceProperty("url", configuration.getUrl());
        dataSource.addDataSourceProperty("user", configuration.getUser());
        dataSource.addDataSourceProperty("password", configuration.getPassword());

        return dataSource;
    }

    @Scheduled(fixedDelay = 5000L)
    public void refreshDataSources() {
        DatabaseConfiguration[] configurations = getDatabaseConfigurations();

        removeObsoleteTenants(configurations);
        insertOrUpdateTenants(configurations);
    }

    private void insertOrUpdateTenants(DatabaseConfiguration[] configurations) {
        for (DatabaseConfiguration configuration : configurations) {
            if (tenants.containsKey(configuration.getTenant())) {
                HikariDataSource dataSource = tenants.get(configuration.getTenant());
                // We only shutdown and reload, if the configuration has actually changed...
                if (!isCurrentConfiguration(dataSource, configuration)) {
                    // Make sure we close this DataSource first...
                    dataSource.close();
                    // ... and then insert a new DataSource:
                    tenants.put(configuration.getTenant(), buildDataSource(configuration));
                }
            } else {
                tenants.put(configuration.getTenant(), buildDataSource(configuration));
            }
        }
    }
    private void removeObsoleteTenants(DatabaseConfiguration[] configurations) {

        // Are there Tenants, that have been removed:
        Set<String> tenantNamesFromConfiguration = Arrays.stream(configurations)
                .map(x -> x.getTenant())
                .collect(Collectors.toSet());

        for (String tenant : tenants.keySet()) {

            // There is currently a Tenant, which is not listed anymore:
            if(!tenantNamesFromConfiguration.contains(tenant)) {

                // So get the DataSource first ...
                HikariDataSource dataSource = tenants.get(tenant);

                // ... close all existing connections:
                dataSource.close();

                // ... and remove it:
                tenants.remove(tenant);
            }
        }
    }

    private boolean isCurrentConfiguration(HikariDataSource dataSource, DatabaseConfiguration configuration) {
        return Objects.equals(dataSource.getDataSourceProperties().getProperty("user"), configuration.getUser())
                && Objects.equals(dataSource.getDataSourceProperties().getProperty("url"), configuration.getUrl())
                && Objects.equals(dataSource.getDataSourceProperties().getProperty("password"), configuration.getPassword())
                && Objects.equals(dataSource.getDataSourceClassName(), configuration.getDataSourceClassName());
    }
}