package cn.lihongjie.coal.spring.config;

import org.hibernate.resource.transaction.spi.DdlTransactionIsolator;
import org.hibernate.tool.schema.internal.HibernateSchemaManagementTool;
import org.hibernate.tool.schema.internal.exec.GenerationTarget;
import org.hibernate.tool.schema.internal.exec.GenerationTargetToDatabase;
import org.hibernate.tool.schema.internal.exec.JdbcContext;
import org.hibernate.tool.schema.spi.ScriptSourceInput;

import java.util.List;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.ForkJoinTask;

public class MultithreadHibernateSchemaManagementTool extends HibernateSchemaManagementTool {

    @Override
    protected GenerationTarget buildDatabaseTarget(
            JdbcContext jdbcContext, boolean needsAutoCommit) {
        return new MTGenerationTargetToDatabase(
                getDdlTransactionIsolator(jdbcContext), true, needsAutoCommit);
    }

    public static class MTGenerationTargetToDatabase extends GenerationTargetToDatabase {

        private DdlTransactionIsolator ddlTransactionIsolator;
        private boolean releaseAfterUse;
        private boolean autocommit;

        private List<ForkJoinTask> tasks;

        public MTGenerationTargetToDatabase(DdlTransactionIsolator ddlTransactionIsolator) {
            super(ddlTransactionIsolator);
        }

        public MTGenerationTargetToDatabase(
                DdlTransactionIsolator ddlTransactionIsolator, boolean releaseAfterUse) {
            super(ddlTransactionIsolator, releaseAfterUse);
        }

        public MTGenerationTargetToDatabase(
                DdlTransactionIsolator ddlTransactionIsolator,
                boolean releaseAfterUse,
                boolean autocommit) {
            super(ddlTransactionIsolator, releaseAfterUse, autocommit);
            this.ddlTransactionIsolator = ddlTransactionIsolator;
            this.releaseAfterUse = releaseAfterUse;
            this.autocommit = autocommit;
        }

        @Override
        public void prepare() {}

        @Override
        public void beforeScript(ScriptSourceInput scriptSource) {
            super.beforeScript(scriptSource);
        }

        @Override
        public void accept(String command) {

            ForkJoinTask<?> task =
                    ForkJoinPool.commonPool()
                            .submit(
                                    () -> {
                                        GenerationTargetToDatabase database =
                                                new GenerationTargetToDatabase(
                                                        ddlTransactionIsolator, true, autocommit);

                                        database.prepare();
                                        database.accept(command);

                                        database.release();
                                    });

            this.tasks.add(task);
        }

        @Override
        public void release() {

            this.tasks.forEach(
                    x -> {
                        try {
                            x.get();
                        } catch (InterruptedException e) {
                            throw new RuntimeException(e);
                        } catch (ExecutionException e) {
                            throw new RuntimeException(e);
                        }
                    });
        }
    }
}
