/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.shardingsphere.core.rewrite.token.generator;

import com.google.common.base.Optional;
import org.apache.shardingsphere.core.parse.sql.context.limit.Limit;
import org.apache.shardingsphere.core.parse.sql.context.limit.LimitValue;
import org.apache.shardingsphere.core.parse.sql.segment.dml.limit.LimitSegment;
import org.apache.shardingsphere.core.parse.sql.segment.dml.limit.LimitValueSegment;
import org.apache.shardingsphere.core.parse.sql.segment.dml.limit.NumberLiteralLimitValueSegment;
import org.apache.shardingsphere.core.parse.sql.statement.SQLStatement;
import org.apache.shardingsphere.core.parse.sql.statement.dml.SelectStatement;
import org.apache.shardingsphere.core.rewrite.token.pojo.OffsetToken;
import org.apache.shardingsphere.core.rule.ShardingRule;

/**
 * Row count token generator.
 *
 * @author panjuan
 */
public final class RowCountTokenGenerator implements OptionalSQLTokenGenerator<ShardingRule> {
    
    @Override
    public Optional<OffsetToken> generateSQLToken(final SQLStatement sqlStatement, final ShardingRule shardingRule) {
        if (!(sqlStatement instanceof SelectStatement)) {
            return Optional.absent();
        }
        if (isExistedOfNumberOffset(sqlStatement)) {
            LimitValueSegment offset = sqlStatement.findSQLSegment(LimitSegment.class).get().getOffset().get();
            return Optional.of(new OffsetToken(offset.getStartIndex(), offset.getStopIndex(), ((NumberLiteralLimitValueSegment) offset).getValue()));
        }
        if (isExistedOfLimitRowNum((SelectStatement) sqlStatement)) {
            LimitValue offset = ((SelectStatement) sqlStatement).getLimit().getOffset();
            return Optional.of(new OffsetToken(offset.getSqlSegment().getStartIndex(), offset.getSqlSegment().getStopIndex(), offset.getValue()));
        }
        return Optional.absent();
    }
    
    private boolean isExistedOfLimitRowNum(final SelectStatement sqlStatement) {
        Limit limit = sqlStatement.getLimit();
        return null != limit && null != limit.getOffset() && -1 != limit.getOffset().getValue();
    }
    
    private boolean isExistedOfNumberOffset(final SQLStatement sqlStatement) {
        Optional<LimitSegment> limitSegment = sqlStatement.findSQLSegment(LimitSegment.class);
        return limitSegment.isPresent() && limitSegment.get().getOffset().isPresent() && limitSegment.get().getOffset().get() instanceof NumberLiteralLimitValueSegment;
    }
}

@@ -1500,32 +1500,36 @@
 ore.
-rewrite.token.pojo.Offse
+parse.sql.token.impl.RowCoun
 tTok
@@ -1784,21 +1784,23 @@
 ptional%3C
-Offse
+RowCoun
 tToken%3E 
@@ -2019,21 +2019,23 @@
 OfNumber
-Offse
+RowCoun
 t(sqlSta
@@ -2075,21 +2075,23 @@
 Segment 
-offse
+rowCoun
 t = sqlS
@@ -2135,37 +2135,39 @@
 class).get().get
-Offse
+RowCoun
 t().get();%0A     
@@ -2196,33 +2196,37 @@
 .of(new 
-OffsetToken(offse
+RowCountToken(rowCoun
 t.getSta
@@ -2232,29 +2232,31 @@
 artIndex(), 
-offse
+rowCoun
 t.getStopInd
@@ -2295,21 +2295,23 @@
 egment) 
-offse
+rowCoun
 t).getVa
@@ -2422,21 +2422,23 @@
 itValue 
-offse
+rowCoun
 t = ((Se
@@ -2480,21 +2480,23 @@
 it().get
-Offse
+RowCoun
 t();%0A   
@@ -2531,25 +2531,29 @@
 new 
-OffsetToken(offse
+RowCountToken(rowCoun
 t.ge
@@ -2579,29 +2579,31 @@
 artIndex(), 
-offse
+rowCoun
 t.getSqlSegm
@@ -2628,13 +2628,15 @@
 (), 
-offse
+rowCoun
 t.ge
@@ -2879,21 +2879,23 @@
 imit.get
-Offse
+RowCoun
 t() && -
@@ -2904,29 +2904,31 @@
 != limit.get
-Offse
+RowCoun
 t().getValue
@@ -2979,21 +2979,23 @@
 OfNumber
-Offse
+RowCoun
 t(final 
@@ -3179,29 +3179,31 @@
 nt.get().get
-Offse
+RowCoun
 t().isPresen
@@ -3235,13 +3235,15 @@
 .get
-Offse
+RowCoun
 t().

modify RowCountTokenGenerator