public class SortTest extends SearchTestCase {

   private static FullTextSession fullTextSession;

   public void setUp() throws Exception {
      super.setUp();
      fullTextSession = Search.getFullTextSession( openSession() );
      createTestNumbers();
   }

   public void tearDown() throws Exception {
      // check for ongoing transaction which is an indicator that something went wrong
      // don't call the cleanup methods in this case. Otherwise the original error get swallowed
      if ( !fullTextSession.getTransaction().isActive() ) {
         deleteTestNumbers();
         fullTextSession.close();
      }
      super.tearDown();
   }

   @SuppressWarnings("unchecked")
   public void testCustomFieldComparatorAscendingSort() {
      Transaction tx = fullTextSession.beginTransaction();

      Query query = new MatchAllDocsQuery();
      FullTextQuery hibQuery = fullTextSession.createFullTextQuery( query, NumberHolder.class );
      Sort sort = new Sort( new SortField( "sum", new SumFieldComparatorSource() ) );
      hibQuery.setSort( sort );
      List<NumberHolder> result = hibQuery.list();
      assertNotNull( result );
      assertEquals( "Wrong number of test results.", 4, result.size() );

      int previousSum = 0;
      for ( NumberHolder n : result ) {
         assertTrue( "Documents should be ordered by increasing sum", previousSum < n.getSum() );
         previousSum = n.getSum();
      }

      tx.commit();
   }

   @SuppressWarnings("unchecked")
   public void testCustomFieldComparatorDescendingSort() {
      Transaction tx = fullTextSession.beginTransaction();

      Query query = new MatchAllDocsQuery();
      FullTextQuery hibQuery = fullTextSession.createFullTextQuery( query, NumberHolder.class );
      Sort sort = new Sort( new SortField( "sum", new SumFieldComparatorSource(), true ) );
      hibQuery.setSort( sort );
      List<NumberHolder> result = hibQuery.list();
      assertNotNull( result );
      assertEquals( "Wrong number of test results.", 4, result.size() );

      int previousSum = 100;
      for ( NumberHolder n : result ) {
         assertTrue( "Documents should be ordered by decreasing sum", previousSum > n.getSum() );
         previousSum = n.getSum();
      }

      tx.commit();
   }

   /**
    * Helper method creating test data for number holder.
    */
   private void createTestNumbers() {
      Transaction tx = fullTextSession.beginTransaction();
      NumberHolder holder = new NumberHolder( 1, 1 );
      fullTextSession.save( holder );
      holder = new NumberHolder( 1, 10 );
      fullTextSession.save( holder );
      holder = new NumberHolder( 1, 5 );
      fullTextSession.save( holder );
      holder = new NumberHolder( 3, 2 );
      fullTextSession.save( holder );
      tx.commit();
      fullTextSession.clear();
   }

   private void deleteTestNumbers() {
      Transaction tx = fullTextSession.beginTransaction();
      fullTextSession.createQuery( "delete " + NumberHolder.class.getName() ).executeUpdate();
      tx.commit();
      fullTextSession.clear();
   }

   protected Class<?>[] getAnnotatedClasses() {
      return new Class[] {
            NumberHolder.class
      };
   }

   @Entity
   @Indexed
   public static class NumberHolder {
      @Id
      @GeneratedValue
      int id;

      @Field(analyze = Analyze.NO)
      int num1;

      @Field(analyze = Analyze.NO)
      int num2;

      public NumberHolder(int num1, int num2) {
         this.num1 = num1;
         this.num2 = num2;
      }

      public int getSum() {
         return num1 + num2;
      }

      @Override
      public String toString() {
         final StringBuilder sb = new StringBuilder();
         sb.append( "NumbersHolder" );
         sb.append( "{id=" ).append( id );
         sb.append( ", num1=" ).append( num1 );
         sb.append( ", num2=" ).append( num2 );
         sb.append( '}' );
         return sb.toString();
      }

      @SuppressWarnings("unused")
      private NumberHolder() {
      }
   }

   public static class SumFieldComparatorSource extends FieldComparatorSource {
      @Override
      public FieldComparator<?> newComparator(String fieldName, int numHits, int sortPos, boolean reversed)
            throws IOException {
         return new SumFieldComparator( numHits, "num1", "num2" );
      }
   }

   public static class SumFieldComparator extends FieldComparator<Integer> {
      private final String field1;
      private final String field2;
      private final int[] field1Values;
      private final int[] field2Values;

      private int[] currentReaderValuesField1;
      private int[] currentReaderValuesField2;
      private int bottom;


      public SumFieldComparator(int numHits, String field1, String field2) {
         this.field1 = field1;
         this.field2 = field2;
         this.field1Values = new int[numHits];
         this.field2Values = new int[numHits];
      }

      @Override
      public int compare(int slot1, int slot2) {
         final int v1 = field1Values[slot1] + field2Values[slot1];
         final int v2 = field1Values[slot2] + field2Values[slot2];

         return compareValues( v1, v2 );
      }


      private int compareValues(int v1, int v2) {
         if ( v1 > v2 ) {
            return 1;
         }
         else if ( v1 < v2 ) {
            return -1;
         }
         else {
            return 0;
         }
      }

      @Override
      public int compareBottom(int doc) {
         int v = currentReaderValuesField1[doc] + currentReaderValuesField2[doc];
         return compareValues( bottom, v );
      }

      @Override
      public void copy(int slot, int doc) {
         int v1 = currentReaderValuesField1[doc];
         field1Values[slot] = v1;

         int v2 = currentReaderValuesField2[doc];
         field2Values[slot] = v2;
      }

      @Override
      public void setNextReader(IndexReader reader, int docBase) throws IOException {
         currentReaderValuesField1 = FieldCache.DEFAULT
               .getInts( reader, field1, FieldCache.DEFAULT_INT_PARSER, false );
         currentReaderValuesField2 = FieldCache.DEFAULT
               .getInts( reader, field2, FieldCache.DEFAULT_INT_PARSER, false );
      }

      @Override
      public void setBottom(final int bottom) {
         this.bottom = field1Values[bottom] + field2Values[bottom];
      }

      @Override
      public Integer value(int slot) {
         return field1Values[slot] + field2Values[slot];
      }
   }
}