package example;
import java.sql.*;
import java.util.concurrent.*;
import javax.sql.DataSource;
import org.springframework.jdbc.datasource.DataSourceUtils;
import org.springframework.jdbc.support.JdbcTransactionManager;
import org.springframework.transaction.support.TransactionTemplate;
import static example.Database.*;
@SuppressWarnings("try")
public final class MultiLab {
    private MultiLab(){}
    static String shardFor(long tenant){
        if(tenant==7)return "A";
        if(tenant==8)return "B";
        throw new IllegalArgumentException("Unknown tenant placement");
    }
    static boolean routeFailure(Throwable failure){
        for(Throwable c=failure;c!=null;c=c.getCause()){
            if(c instanceof IllegalStateException)return true;
        }return false;
    }
    static String database(DataSource source){
        var c=DataSourceUtils.getConnection(source);
        try{return scalar(c,"SELECT current_database()");}
        catch(SQLException failure){throw new IllegalStateException(failure);}
        finally{DataSourceUtils.releaseConnection(c,source);}
    }
    static void first(Routing source)throws Exception{
        try(var a=app(a());var b=app(b())){
            sql(a,"DELETE FROM tenant_order");sql(b,"DELETE FROM tenant_order");
            try{
                try(var route=RouteScope.use("A")){
                    var tx=new TransactionTemplate(new JdbcTransactionManager(source));
                    tx.executeWithoutResult(status->{
                        var c=DataSourceUtils.getConnection(source);
                        try{check("route_a".equals(scalar(c,"SELECT current_database()")),"A identity");sql(c,"INSERT INTO tenant_order VALUES(1,7,'tenant-7',NULL)");}
                        catch(SQLException failure){throw new IllegalStateException(failure);}
                        finally{DataSourceUtils.releaseConnection(c,source);}
                    });
                }
                check("1".equals(scalar(a,"SELECT count(*) FROM tenant_order"))&&"0".equals(scalar(b,"SELECT count(*) FROM tenant_order")),"different databases");
                System.out.println("routeA=route_a routeB=route_b writtenA=1 rowsB=0");
            }finally{sql(a,"DELETE FROM tenant_order");sql(b,"DELETE FROM tenant_order");}
        }
    }
    static void bound(Routing source){
        try(var route=RouteScope.use("A")){
            var tx=new TransactionTemplate(new JdbcTransactionManager(source));
            tx.executeWithoutResult(status->{
                check(database(source).equals("route_a"),"initial A");
                try(var inner=RouteScope.use("B")){check(database(source).equals("route_a"),"same transaction still A");}
            });
        }
        try(var route=RouteScope.use("B")){check(database(source).equals("route_b"),"new work B");}
        System.out.println("transactionBoundToA=true changedRouteStillA=true nextWorkSelectedB=true");
    }
    static void scopes(Routing source)throws Exception{
        try(var route=RouteScope.use("A")){
            try(var inner=RouteScope.use("B")){check(database(source).equals("route_b"),"nested B");}
            check(database(source).equals("route_a"),"outer restored");
            try(var failing=RouteScope.use("B")){throw new IllegalArgumentException("planned scope failure");}
            catch(IllegalArgumentException expected){check(database(source).equals("route_a"),"exception restored outer route");}
            ExecutorService worker=Executors.newSingleThreadExecutor();
            try{
                check(worker.submit(()->{try{RouteScope.required();return false;}catch(IllegalStateException expected){return true;}}).get(5,TimeUnit.SECONDS),"not inherited");
                check(worker.submit(()->{try(var explicit=RouteScope.use("B")){return database(source).equals("route_b");}}).get(5,TimeUnit.SECONDS),"explicit route");
                check(worker.submit(()->{try{RouteScope.required();return false;}catch(IllegalStateException expected){return true;}}).get(5,TimeUnit.SECONDS),"worker cleaned");
            }finally{worker.shutdownNow();check(worker.awaitTermination(5,TimeUnit.SECONDS),"worker stopped");}
        }
        boolean missing=false,unknown=false;
        try{database(source);}catch(RuntimeException expected){missing=routeFailure(expected);}
        try(var route=RouteScope.use("unknown")){try{database(source);}catch(RuntimeException expected){unknown=routeFailure(expected);}}
        check(missing&&unknown,"strict missing and unknown");
        System.out.println("nestedRouteRestored=true exceptionRestoredRoute=true missingRouteRejected=true unknownRouteRejected=true asyncExplicitRoute=true workerContextCleared=true");
    }
    static void shard(Routing source)throws Exception{
        try{
            for(long tenant:new long[]{7,8}){
                String key=shardFor(tenant);
                try(var route=RouteScope.use(key)){
                    var tx=new TransactionTemplate(new JdbcTransactionManager(source));
                    tx.executeWithoutResult(status->{
                        var c=DataSourceUtils.getConnection(source);
                        try(var s=c.prepareStatement("INSERT INTO tenant_order(id,tenant_id,legacy_label) VALUES(?,?,?)")){
                            s.setLong(1,tenant);s.setLong(2,tenant);s.setString(3,"tenant-"+tenant);s.executeUpdate();
                        }catch(SQLException failure){throw new IllegalStateException(failure);}
                        finally{DataSourceUtils.releaseConnection(c,source);}
                    });
                }
            }
            try(var a=app(a());var b=app(b())){
                check("7".equals(scalar(a,"SELECT tenant_id FROM tenant_order"))&&"8".equals(scalar(b,"SELECT tenant_id FROM tenant_order")),"actual shards");
            }
            System.out.println("tenant7StoredA=true tenant8StoredB=true explicitPlacement=true");
        }finally{try(var a=app(a());var b=app(b())){sql(a,"DELETE FROM tenant_order");sql(b,"DELETE FROM tenant_order");}}
    }
    public static void main(String[] args)throws Exception{
        if("root".equals(System.getProperty("user.name")))throw new IllegalStateException("Use non-root runtime");
        Migrations.migrate();String mode=args.length==0?"all":args[0];
        try(var source=new Routing()){
            switch(mode){
                case "first"->first(source);case "bound"->bound(source);case "scopes"->scopes(source);case "shard"->shard(source);
                case "migrate"->Migrations.validate();case "checksum"->Migrations.checksum();case "ddl"->Migrations.ddl();
                case "permissions"->Migrations.permissions();case "backfill"->Migrations.backfill();case "replication"->Replication.run();
                case "all"->{first(source);bound(source);scopes(source);shard(source);Migrations.validate();Migrations.checksum();Migrations.ddl();Migrations.permissions();Migrations.backfill();Replication.run();}
                default->throw new IllegalArgumentException("Unknown mode: "+mode);
            }
        }
    }
}