package dialect

import (
	"fmt"
	"strings"
	"time"

	"github.com/pingcap/tidb/pkg/parser"
	"github.com/pingcap/tidb/pkg/parser/ast"
	"github.com/pingcap/tidb/pkg/parser/opcode"
)

// DateScope is a finite, half-open calendar-date interval derived from SQL.
// End is the first date that is not included in the interval.
type DateScope struct {
	Start time.Time
	End   time.Time
}

// ExtractLiteralDateScope returns a conservative source-planning interval for
// a literal date predicate. The original SQL remains responsible for exact row
// filtering (for example, a non-contiguous IN list).
func ExtractLiteralDateScope(sql, column string, maxDays int) (DateScope, error) {
	p := parser.New()
	stmt, err := p.ParseOneStmt(sql, "", "")
	if err != nil {
		return DateScope{}, fmt.Errorf("parse statement: %w", err)
	}
	if !isSupportedSelectLike(stmt) {
		return DateScope{}, fmt.Errorf("only SELECT and WITH (CTE) queries are supported")
	}

	var whereExprs []ast.ExprNode
	collectWhereExprs(stmt, &whereExprs)
	var result dateInterval
	for _, where := range whereExprs {
		interval, referenced, err := extractDateInterval(where, strings.ToLower(column))
		if err != nil {
			return DateScope{}, err
		}
		if !referenced {
			continue
		}
		if !interval.bounded() {
			return DateScope{}, fmt.Errorf("date partition filter must have finite literal lower and upper bounds")
		}
		result = unionDateIntervals(result, interval)
	}
	if !result.bounded() {
		return DateScope{}, fmt.Errorf("date partition filter must have finite literal lower and upper bounds")
	}
	if !result.lower.Before(result.upper) {
		return DateScope{}, fmt.Errorf("date partition range is empty or reversed")
	}
	if maxDays > 0 && int(result.upper.Sub(result.lower)/(24*time.Hour)) > maxDays {
		return DateScope{}, fmt.Errorf("date partition range exceeds %d days", maxDays)
	}
	return DateScope{Start: result.lower, End: result.upper}, nil
}

type dateInterval struct {
	lower time.Time
	upper time.Time
}

func (i dateInterval) bounded() bool {
	return !i.lower.IsZero() && !i.upper.IsZero()
}

func extractDateInterval(expr ast.ExprNode, column string) (dateInterval, bool, error) {
	if expr == nil {
		return dateInterval{}, false, nil
	}
	expr = peelExpr(expr)
	switch n := expr.(type) {
	case *ast.BinaryOperationExpr:
		if n.Op == opcode.LogicAnd || n.Op == opcode.LogicOr {
			left, leftReferenced, err := extractDateInterval(n.L, column)
			if err != nil {
				return dateInterval{}, true, err
			}
			right, rightReferenced, err := extractDateInterval(n.R, column)
			if err != nil {
				return dateInterval{}, true, err
			}
			if n.Op == opcode.LogicOr {
				if leftReferenced != rightReferenced || (leftReferenced && (!left.bounded() || !right.bounded())) {
					return dateInterval{}, true, fmt.Errorf("date partition filter contains an unbounded OR branch")
				}
				if !leftReferenced {
					return dateInterval{}, false, nil
				}
				return unionDateIntervals(left, right), true, nil
			}
			if !leftReferenced {
				return right, rightReferenced, nil
			}
			if !rightReferenced {
				return left, true, nil
			}
			return intersectDateIntervals(left, right), true, nil
		}

		leftIsDate := isPartitionColRef(n.L, column)
		rightIsDate := isPartitionColRef(n.R, column)
		if !leftIsDate && !rightIsDate {
			return dateInterval{}, false, nil
		}
		if leftIsDate == rightIsDate {
			return dateInterval{}, true, fmt.Errorf("date partition filter must compare the date column with a literal")
		}
		literal := n.R
		op := n.Op
		if rightIsDate {
			literal = n.L
			op = reverseComparison(op)
		}
		date, err := literalDate(literal)
		if err != nil {
			return dateInterval{}, true, err
		}
		switch op {
		case opcode.EQ:
			return dateInterval{lower: date, upper: date.AddDate(0, 0, 1)}, true, nil
		case opcode.GE:
			return dateInterval{lower: date}, true, nil
		case opcode.GT:
			return dateInterval{lower: date.AddDate(0, 0, 1)}, true, nil
		case opcode.LE:
			return dateInterval{upper: date.AddDate(0, 0, 1)}, true, nil
		case opcode.LT:
			return dateInterval{upper: date}, true, nil
		default:
			return dateInterval{}, true, fmt.Errorf("unsupported date partition comparison")
		}
	case *ast.PatternInExpr:
		if !isPartitionColRef(n.Expr, column) {
			return dateInterval{}, false, nil
		}
		if n.Not || n.Sel != nil || len(n.List) == 0 {
			return dateInterval{}, true, fmt.Errorf("date IN filter must contain literal values")
		}
		var interval dateInterval
		for _, item := range n.List {
			date, err := literalDate(item)
			if err != nil {
				return dateInterval{}, true, err
			}
			interval = unionDateIntervals(interval, dateInterval{lower: date, upper: date.AddDate(0, 0, 1)})
		}
		return interval, true, nil
	case *ast.BetweenExpr:
		if !isPartitionColRef(n.Expr, column) {
			return dateInterval{}, false, nil
		}
		if n.Not {
			return dateInterval{}, true, fmt.Errorf("NOT BETWEEN is not supported for date source planning")
		}
		lower, err := literalDate(n.Left)
		if err != nil {
			return dateInterval{}, true, err
		}
		upper, err := literalDate(n.Right)
		if err != nil {
			return dateInterval{}, true, err
		}
		return dateInterval{lower: lower, upper: upper.AddDate(0, 0, 1)}, true, nil
	default:
		return dateInterval{}, false, nil
	}
}

func literalDate(expr ast.ExprNode) (time.Time, error) {
	value, ok := peelExpr(expr).(ast.ValueExpr)
	if !ok {
		return time.Time{}, fmt.Errorf("date partition bounds must be literal YYYY-MM-DD values")
	}
	text, ok := value.GetValue().(string)
	if !ok {
		return time.Time{}, fmt.Errorf("date partition bounds must be literal YYYY-MM-DD values")
	}
	date, err := time.Parse("2006-01-02", text)
	if err != nil {
		return time.Time{}, fmt.Errorf("invalid date partition literal %q; expected YYYY-MM-DD", text)
	}
	return date.UTC(), nil
}

func reverseComparison(op opcode.Op) opcode.Op {
	switch op {
	case opcode.GE:
		return opcode.LE
	case opcode.GT:
		return opcode.LT
	case opcode.LE:
		return opcode.GE
	case opcode.LT:
		return opcode.GT
	default:
		return op
	}
}

func unionDateIntervals(left, right dateInterval) dateInterval {
	if left.lower.IsZero() {
		left.lower = right.lower
	} else if !right.lower.IsZero() && right.lower.Before(left.lower) {
		left.lower = right.lower
	}
	if left.upper.IsZero() || (!right.upper.IsZero() && right.upper.After(left.upper)) {
		left.upper = right.upper
	}
	return left
}

func intersectDateIntervals(left, right dateInterval) dateInterval {
	result := left
	if result.lower.IsZero() || (!right.lower.IsZero() && right.lower.After(result.lower)) {
		result.lower = right.lower
	}
	if result.upper.IsZero() || (!right.upper.IsZero() && right.upper.Before(result.upper)) {
		result.upper = right.upper
	}
	return result
}