package example;
import jakarta.persistence.*;
import java.sql.*;
import java.util.*;
import java.util.concurrent.*;
import org.hibernate.SessionFactory;

public final class PageLab {
    private PageLab(){}
    static void check(boolean ok,String text){if(!ok)throw new AssertionError(text);}
    static Connection connect() throws SQLException {return Factories.dataSource().getConnection();}
    static void execute(String sql)throws SQLException{
        try(var c=connect();var s=c.createStatement()){s.executeUpdate(sql);}
    }
    static List<Long> ids(Connection c,String sql,long... parameters)throws SQLException{
        try(var s=c.prepareStatement(sql)){
            s.setQueryTimeout(5);
            for(int i=0;i<parameters.length;i++)s.setLong(i+1,parameters[i]);
            try(var r=s.executeQuery()){var out=new ArrayList<Long>();while(r.next())out.add(r.getLong(1));return out;}
        }
    }
    static List<Long> offset(long value)throws SQLException{
        try(var c=connect()){return ids(c,"SELECT id FROM catalog_item ORDER BY sort_seq,id LIMIT 2 OFFSET ?",value);}
    }
    static List<Long> keyset(long seq,long id)throws SQLException{
        try(var c=connect()){return ids(c,"SELECT id FROM catalog_item WHERE (sort_seq,id)>(?,?) ORDER BY sort_seq,id LIMIT 2",seq,id);}
    }
    static void first()throws SQLException{
        check(offset(0).equals(List.of(1L,2L)),"first");check(offset(2).equals(List.of(3L,4L)),"second");
        check(keyset(20,2).equals(List.of(3L,4L)),"keyset");
        System.out.println("firstPage=[1, 2] offsetSecond=[3, 4] keysetSecond=[3, 4]");
    }
    static void drift()throws SQLException{
        try{
            var first=offset(0);execute("INSERT INTO catalog_item VALUES(99,5,10,7)");
            var offset=offset(2);var keyset=keyset(20,2);
            check(first.equals(List.of(1L,2L))&&offset.equals(List.of(2L,3L))&&keyset.equals(List.of(3L,4L)),"insert drift");
            System.out.println("insertBeforePageOffset=[2, 3] keyset=[3, 4] offsetRepeatedId=true");
            execute("UPDATE catalog_item SET sort_seq=25 WHERE id=1");
            check(keyset(20,2).equals(List.of(3L,1L)),"moved sort key repeats");
            System.out.println("sortKeyChangedKeyset=[3, 1] previouslySeenIdRepeated=true");
        }finally{execute("DELETE FROM catalog_item WHERE id=99");execute("UPDATE catalog_item SET sort_seq=10 WHERE id=1");}
    }
    static Map<Long,List<String>> graph(List<CatalogItem> items){
        var out=new LinkedHashMap<Long,List<String>>();
        for(var item:items)out.put(item.getId(),item.getLines().stream().map(CatalogLine::getSku).sorted().toList());
        return out;
    }
    static void nplus() {
        try(var f=Factories.create()){
            var stats=f.unwrap(SessionFactory.class).getStatistics();stats.clear();
            Map<Long,List<String>> slow;
            try(var em=f.createEntityManager()){
                var items=em.createQuery("select i from CatalogItem i order by i.sortSeq,i.id",CatalogItem.class).setMaxResults(3).getResultList();
                slow=graph(items);
            }
            long slowCount=stats.getPrepareStatementCount();check(slowCount==4,"N+1 actual statements");
            stats.clear();Map<Long,List<String>> fast;
            try(var em=f.createEntityManager()){
                var page=em.createQuery("select i.id from CatalogItem i order by i.sortSeq,i.id",Long.class).setMaxResults(3).getResultList();
                var items=em.createQuery("select i from CatalogItem i left join fetch i.lines where i.id in :ids",CatalogItem.class).setParameter("ids",page).getResultList();
                var unordered=graph(items);fast=new LinkedHashMap<>();for(Long id:page)fast.put(id,unordered.get(id));
            }
            check(stats.getPrepareStatementCount()==2&&slow.equals(fast)&&new ArrayList<>(fast.keySet()).equals(List.of(1L,2L,3L)),"two queries same graph/order");
            System.out.println("nPlusOneStatements=4 twoStepStatements=2 graphEqual=true orderedIds=[1, 2, 3]");
            stats.clear();SqlTrace.SQL.clear();
            try(var em=f.createEntityManager()){
                var fetched=em.createQuery("select i from CatalogItem i left join fetch i.lines order by i.sortSeq,i.id",CatalogItem.class).setMaxResults(3).getResultList();
                check(graph(fetched).equals(fast)&&fetched.size()==3,"collection fetch page graph");
            }
            check(stats.getPrepareStatementCount()==1&&SqlTrace.SQL.size()==1,"one fetch page statement");
            System.out.println("fetchPageSql="+SqlTrace.SQL.get(0));
            String generated=SqlTrace.SQL.get(0).toLowerCase(Locale.ROOT);
            check(generated.contains("fetch first")&&generated.contains("(select"),"database subquery limit");
            System.out.println("hibernate74DatabaseFetchPage=true statements=1 parents=3 childRows=6 completeGraph=true");
        }
    }
    static void plan()throws SQLException{
        check(offset(10000).equals(keyset(10994,10994)),"plan queries return same page");
        for(String sql:List.of(
            "SELECT id FROM catalog_item ORDER BY sort_seq,id LIMIT 2 OFFSET 10000",
            "SELECT id FROM catalog_item WHERE (sort_seq,id)>(10994,10994) ORDER BY sort_seq,id LIMIT 2")){
            try(var c=connect();var s=c.createStatement();var r=s.executeQuery("EXPLAIN (ANALYZE,BUFFERS) "+sql)){
                int lines=0;while(r.next()){System.out.println(r.getString(1));lines++;}
                check(lines>1,"actual plan");
            }
        }
        System.out.println("offsetPlanProduced=true keysetPlanProduced=true");
    }
    static long version(Connection c)throws SQLException{return ids(c,"SELECT version FROM catalog_item WHERE id=1").get(0);}
    static int debit(Connection c,long v)throws SQLException{
        try(var s=c.prepareStatement("UPDATE catalog_item SET quantity=quantity-1,version=version+1 WHERE id=1 AND version=? AND quantity>0")){
            s.setLong(1,v);return s.executeUpdate();
        }
    }
    static void optimistic()throws SQLException{
        try(var a=connect();var b=connect()){
            a.setAutoCommit(false);b.setAutoCommit(false);
            try{
                long av=version(a),bv=version(b);check(av==7&&bv==7,"same initial version");
                check(debit(a,av)==1,"first update");a.commit();
                check(debit(b,bv)==0,"old version");b.rollback();
                System.out.println("initialVersion=7 firstUpdateRows=1 staleUpdateRows=0");
            }finally{a.rollback();b.rollback();execute("UPDATE catalog_item SET quantity=10,version=7 WHERE id=1");}
        }
    }
    static void nowait()throws SQLException{
        try(var a=connect();var b=connect()){
            a.setAutoCommit(false);b.setAutoCommit(false);
            try{
                ids(a,"SELECT id FROM catalog_item WHERE id=1 FOR UPDATE");
                boolean rejected=false;
                try{ids(b,"SELECT id FROM catalog_item WHERE id=1 FOR UPDATE NOWAIT");}
                catch(SQLException expected){rejected="55P03".equals(expected.getSQLState());}
                check(rejected,"nowait rejected");b.rollback();a.commit();
                check(ids(b,"SELECT id FROM catalog_item WHERE id=1 FOR UPDATE NOWAIT").equals(List.of(1L)),"after release");
                System.out.println("nowaitRejected=true sqlState=55P03 lockAfterRelease=true");
            }finally{a.rollback();b.rollback();}
        }
    }
    static void skipLocked()throws SQLException{
        try(var a=connect();var b=connect()){
            a.setAutoCommit(false);b.setAutoCommit(false);
            try{
                String sql="SELECT id FROM task WHERE state='ready' ORDER BY id LIMIT 1 FOR UPDATE SKIP LOCKED";
                var first=ids(a,sql);var second=ids(b,sql);
                check(first.equals(List.of(1L))&&second.equals(List.of(2L)),"skip locked work item");
                System.out.println("consumerA=1 consumerB=2 differentLockedTasks=true");
            }finally{a.rollback();b.rollback();}
        }
    }
    static void lockWait()throws Exception{
        ExecutorService worker=Executors.newSingleThreadExecutor();
        try(var a=connect();var b=connect();var observer=connect()){
            a.setAutoCommit(false);b.setAutoCommit(false);
            try{
                ids(a,"SELECT id FROM catalog_item WHERE id=1 FOR UPDATE");
                long bp=ids(b,"SELECT pg_backend_pid()").get(0);
                Future<List<Long>> waiting=worker.submit(()->ids(b,"SELECT id FROM catalog_item WHERE id=1 FOR UPDATE"));
                long deadline=System.nanoTime()+TimeUnit.SECONDS.toNanos(3);boolean observed=false;
                do{
                    var blocked=ids(observer,"SELECT count(*) FROM pg_stat_activity WHERE pid=? AND wait_event_type='Lock' AND cardinality(pg_blocking_pids(pid))>0",bp);
                    observed=blocked.get(0)==1;
                    if(!observed)Thread.sleep(10);
                }while(!observed&&System.nanoTime()<deadline);
                check(observed,"database wait evidence");a.commit();
                check(waiting.get(5,TimeUnit.SECONDS).equals(List.of(1L)),"released waiter");
                System.out.println("databaseLockWaitObserved=true blockerPidLinked=true releaseCompletedQuery=true");
            }finally{a.rollback();b.rollback();}
        }finally{worker.shutdownNow();check(worker.awaitTermination(10,TimeUnit.SECONDS),"worker stopped");}
    }
    public static void main(String[] args)throws Exception{
        if("root".equals(System.getProperty("user.name")))throw new IllegalStateException("Use non-root runtime");
        String mode=args.length==0?"all":args[0];
        switch(mode){
            case "first"->first();case "drift"->drift();case "nplus"->nplus();case "plan"->plan();
            case "optimistic"->optimistic();case "nowait"->nowait();case "skip-locked"->skipLocked();case "lock-wait"->lockWait();
            case "all"->{first();drift();nplus();plan();optimistic();nowait();skipLocked();lockWait();}
            default->throw new IllegalArgumentException("Unknown mode: "+mode);
        }
    }
}