TsDB2/DB/core/query.go

package core

import (
	"TsDB2/DB/core/index"
	"TsDB2/DB/core/page"
	"TsDB2/DB/core/types"
	"bytes"
	"encoding/binary"
	"sort"
)

// Query 查询操作对象
type Query struct {
	table      *Table
	selectCols []string    // SELECT 字段列表
	whereCond  Condition   // WHERE 条件
	orderBys   []SortField // ORDER BY
	limit      int         // LIMIT
	offset     int         // OFFSET
	groupBys   []string    // GROUP BY
	havingCond Condition   // HAVING
}

type SortField struct {
	Field      string
	Descending bool
}

// scanPlan 索引扫描计划
type scanPlan struct {
	meta           *IndexMeta
	handler        index.IndexHandler
	startKey       []byte
	endKey         []byte
	startInclusive bool
	endInclusive   bool
	isPrimary      bool      // 主键/时序索引=true
	remainingCond  Condition // 需要行级过滤的剩余条件
}

// Query 返回查询操作对象
func (table *Table) Query() *Query {
	return &Query{table: table}
}

func (q *Query) Select(cols ...string) *Query  { q.selectCols = cols; return q }
func (q *Query) Where(cond Condition) *Query   { q.whereCond = cond; return q }
func (q *Query) Limit(n int) *Query            { q.limit = n; return q }
func (q *Query) Offset(n int) *Query           { q.offset = n; return q }
func (q *Query) GroupBy(cols ...string) *Query { q.groupBys = cols; return q }
func (q *Query) Having(cond Condition) *Query  { q.havingCond = cond; return q }

func (q *Query) OrderBy(field string, descending ...bool) *Query {
	desc := false
	if len(descending) > 0 {
		desc = descending[0]
	}
	q.orderBys = append(q.orderBys, SortField{Field: field, Descending: desc})
	return q
}

// ======================== 索引键编码 ========================

// encodeIndexKey 将条件值编码为索引键字节 (BigEndian, 与 writeFieldValue 一致)
func encodeIndexKey(value interface{}, ft types.FieldType) []byte {
	size := types.GetSize(ft)
	buf := make([]byte, size)
	switch ft {
	case types.FieldType_1byte:
		buf[0] = value.(byte)
	case types.FieldType_2byte:
		binary.BigEndian.PutUint16(buf, uint16(value.(int16)))
	case types.FieldType_4byte:
		binary.BigEndian.PutUint32(buf, uint32(value.(int32)))
	case types.FieldType_8byte:
		binary.BigEndian.PutUint64(buf, value.(uint64))
	default:
		// 字符串类型
		b := value.([]byte)
		copy(buf, b)
	}
	return buf
}

// ======================== 索引匹配 ========================

// fieldNameFromMeta 获取索引第一个字段的名称
func (q *Query) fieldNameFromMeta(meta *IndexMeta) string {
	fieldIdx := int(meta.FieldIDs[0])
	if fieldIdx < len(q.table.Field) {
		return string(q.table.Field[fieldIdx].FieldName)
	}
	return ""
}

// fieldTypeFromMeta 获取索引第一个字段的类型
func (q *Query) fieldTypeFromMeta(meta *IndexMeta) types.FieldType {
	fieldIdx := int(meta.FieldIDs[0])
	if fieldIdx < len(q.table.Field) {
		return q.table.Field[fieldIdx].FieldType
	}
	return 0
}

// condToRange 尝试从条件中提取扫描范围, 返回 (startKey, endKey, startInc, endInc, ok)
func (q *Query) condToRange(cond Condition, fieldName string, ft types.FieldType) ([]byte, []byte, bool, bool, Condition) {
	switch c := cond.(type) {
	case EqCondition:
		if c.Field == fieldName {
			key := encodeIndexKey(c.Value, ft)
			return key, key, true, true, nil
		}
	case GeCondition:
		if c.Field == fieldName {
			return encodeIndexKey(c.Value, ft), nil, true, false, nil
		}
	case GtCondition:
		if c.Field == fieldName {
			return encodeIndexKey(c.Value, ft), nil, false, false, nil
		}
	case LeCondition:
		if c.Field == fieldName {
			return nil, encodeIndexKey(c.Value, ft), false, true, nil
		}
	case LtCondition:
		if c.Field == fieldName {
			return nil, encodeIndexKey(c.Value, ft), false, false, nil
		}
	case AndCondition:
		// 合并所有子条件的范围
		var sk, ek []byte
		var si, ei bool = true, true // 默认全包含
		var hasStart, hasEnd bool
		var remaining []Condition

		for _, sub := range c.Conditions {
			s, e, sinc, einc, rem2 := q.condToRange(sub, fieldName, ft)
			if rem2 == nil && (s != nil || e != nil) {
				// 此子条件可完全用于索引
				if s != nil {
					if !hasStart || (sinc && bytes.Compare(s, sk) > 0) || (!sinc && bytes.Compare(s, sk) >= 0) {
						sk = s
						si = sinc
						hasStart = true
					}
				}
				if e != nil {
					if !hasEnd || (einc && bytes.Compare(e, ek) < 0) || (!einc && bytes.Compare(e, ek) <= 0) {
						ek = e
						ei = einc
						hasEnd = true
					}
				}
			} else if s != nil || e != nil {
				// 此子条件部分可用于索引
				if s != nil {
					if !hasStart || (sinc && bytes.Compare(s, sk) > 0) || (!sinc && bytes.Compare(s, sk) >= 0) {
						sk = s
						si = sinc
						hasStart = true
					}
				}
				if e != nil {
					if !hasEnd || (einc && bytes.Compare(e, ek) < 0) || (!einc && bytes.Compare(e, ek) <= 0) {
						ek = e
						ei = einc
						hasEnd = true
					}
				}
				if rem2 != nil {
					remaining = append(remaining, rem2)
				}
			} else {
				// 此子条件不能用于索引
				remaining = append(remaining, sub)
			}
		}

		if hasStart || hasEnd {
			var rem Condition
			if len(remaining) == 1 {
				rem = remaining[0]
			} else if len(remaining) > 1 {
				rem = And(remaining...)
			}
			return sk, ek, si, ei, rem
		}
	}
	return nil, nil, false, false, cond
}

// buildScanPlan 分析 WHERE 条件, 选择最优索引
func (q *Query) buildScanPlan() *scanPlan {
	if q.whereCond == nil {
		return nil
	}
	if q.table.indexMeta == nil {
		return nil
	}

	var bestPlan *scanPlan

	for i := range q.table.indexMeta.metas {
		m := &q.table.indexMeta.metas[i]
		if m.Flag&0x80 != 0 {
			continue
		}
		isPrimary := m.IndexType == page.IndexTypePrimary || m.IndexType == page.IndexTypeTime
		fieldName := q.fieldNameFromMeta(m)
		ft := q.fieldTypeFromMeta(m)
		if fieldName == "" {
			continue
		}

		sk, ek, si, ei, rem := q.condToRange(q.whereCond, fieldName, ft)
		if sk == nil && ek == nil {
			continue // 此索引不匹配
		}

		handler, ok := index.Compose.GetHandler(m.IndexType)
		if !ok {
			continue
		}

		plan := &scanPlan{
			meta: m, handler: handler,
			startKey: sk, endKey: ek,
			startInclusive: si, endInclusive: ei,
			isPrimary: isPrimary, remainingCond: rem,
		}
		// 优先选主键/时序索引 (避免回查)
		if isPrimary || bestPlan == nil {
			bestPlan = plan
		}
		if isPrimary {
			break // 主键索引最优, 不再找
		}
	}
	return bestPlan
}

// ======================== 核心查询 ========================

// Find 执行查询，返回多行
func (q *Query) Find() ([][]types.FieldValue, error) {
	var results [][]types.FieldValue
	err := q.scan(func(row []types.FieldValue) bool {
		results = append(results, row)
		return true
	})
	if err != nil {
		return nil, err
	}
	// 投影
	if len(q.selectCols) > 0 {
		results = q.project(results)
	}
	return results, nil
}

// Scan 流式扫描查询, 回调函数逐行处理数据, 返回 false 提前终止
// 适用于不需要全量收集结果的场景, 内存友好
func (q *Query) Scan(callback func(row []types.FieldValue) bool) error {
	return q.scan(callback)
}

// scan 内部流式扫描核心方法
// 自动处理: 索引选择 → 排序 → OFFSET → LIMIT → 回调
func (q *Query) scan(callback func(row []types.FieldValue) bool) error {
	plan := q.buildScanPlan()

	// 判断是否需要内存排序 (OrderBy 字段与索引不一致时)
	needSort := false
	if len(q.orderBys) > 0 {
		if plan == nil {
			needSort = true
		} else {
			idxField := q.fieldNameFromMeta(plan.meta)
			if !(len(q.orderBys) == 1 && q.orderBys[0].Field == idxField) {
				needSort = true
			}
		}
	}

	// 需要排序时, 必须先收集全部再排序, 无法流式 LIMIT
	if needSort {
		var all [][]types.FieldValue
		err := q.executeScan(func(row []types.FieldValue) bool {
			all = append(all, row)
			return true
		}, plan)
		if err != nil {
			return err
		}
		q.sortResults(all)
		// OFFSET
		if q.offset > 0 {
			if q.offset >= len(all) {
				return nil
			}
			all = all[q.offset:]
		}
		// LIMIT
		if q.limit > 0 && q.limit < len(all) {
			all = all[:q.limit]
		}
		for _, row := range all {
			if !callback(row) {
				return nil
			}
		}
		return nil
	}

	// 不需要排序: 直接流式传递 callback, 支持 LIMIT 提前终止
	seen := 0
	return q.executeScan(func(row []types.FieldValue) bool {
		if seen < q.offset {
			seen++
			return true // 跳过 OFFSET
		}
		seen++
		if !callback(row) {
			return false // 用户要求停止
		}
		// LIMIT 检查 (0 表示无限制)
		if q.limit > 0 && seen-q.offset >= q.limit {
			return false // 已满足, 提前终止扫描
		}
		return true
	}, plan)
}

// executeScan 执行索引扫描或全表扫描, 将反序列化后的行传递给回调
func (q *Query) executeScan(callback func(row []types.FieldValue) bool, plan *scanPlan) error {
	if plan != nil {
		return q.indexScan(plan, callback)
	}
	return q.fullTableScan(callback)
}

// indexScan 使用索引扫描 (流式回调模式)
func (q *Query) indexScan(plan *scanPlan, callback func(row []types.FieldValue) bool) error {
	// 判断 OrderBy 是否需要反向
	reverse := false
	if len(q.orderBys) == 1 {
		idxField := q.fieldNameFromMeta(plan.meta)
		if q.orderBys[0].Field == idxField && q.orderBys[0].Descending {
			reverse = true
		}
	}

	if plan.isPrimary {
		// 主键/时序索引: RangeScan 直接返回 storeRow = [pkID(8B) + rowData]
		return plan.handler.RangeScan(q.table.handle, plan.meta.EntryPage,
			plan.startKey, plan.endKey, plan.startInclusive, plan.endInclusive, reverse,
			func(row []byte) bool {
				if len(row) < 8 {
					return true
				}
				rowData := row[8:]
				if plan.remainingCond != nil {
					if !plan.remainingCond.Match(rowData, q.table) {
						return true
					}
				}
				return callback(q.table.deserializeRow(rowData))
			})
	}

	// 二级索引: RangeScan 返回 pkID(8B BE), 需回查主键
	return plan.handler.RangeScan(q.table.handle, plan.meta.EntryPage,
		plan.startKey, plan.endKey, plan.startInclusive, plan.endInclusive, reverse,
		func(row []byte) bool {
			if len(row) < 8 {
				return true
			}
			pkID := binary.BigEndian.Uint64(row[:8])
			storeRow, found := q.table.indexMeta.ResolveByPK(pkID, q.table.Field)
			if !found {
				return true
			}
			rowData := storeRow[8:]
			if plan.remainingCond != nil {
				if !plan.remainingCond.Match(rowData, q.table) {
					return true
				}
			}
			return callback(q.table.deserializeRow(rowData))
		})
}

// fullTableScan 全表扫描 (流式回调模式)
func (q *Query) fullTableScan(callback func(row []types.FieldValue) bool) error {
	// 找主键/时序索引
	var primMeta *IndexMeta
	var primHandler index.IndexHandler
	for i := range q.table.indexMeta.metas {
		m := &q.table.indexMeta.metas[i]
		if m.Flag&0x80 != 0 {
			continue
		}
		if m.IndexType == page.IndexTypePrimary || m.IndexType == page.IndexTypeTime {
			primMeta = m
			primHandler, _ = index.Compose.GetHandler(m.IndexType)
			break
		}
	}
	if primMeta == nil || primHandler == nil {
		return nil
	}

	dpIDs := primHandler.GetDataPageIDs(q.table.handle, primMeta.EntryPage)
	for _, dpID := range dpIDs {
		rows := primHandler.ReadDataPageRows(q.table.handle, dpID)
		for _, storeRow := range rows {
			if len(storeRow) < 8 {
				continue
			}
			rowData := storeRow[8:]
			if q.whereCond != nil {
				if !q.whereCond.Match(rowData, q.table) {
					continue
				}
			}
			if !callback(q.table.deserializeRow(rowData)) {
				return nil // 回调要求提前终止
			}
		}
	}
	return nil
}

// ======================== 排序 ========================

func (q *Query) sortResults(results [][]types.FieldValue) {
	sort.SliceStable(results, func(i, j int) bool {
		for _, ob := range q.orderBys {
			fi := q.table.findFieldIndex(ob.Field)
			if fi < 0 {
				continue
			}
			vi := findFieldValue(results[i], ob.Field)
			vj := findFieldValue(results[j], ob.Field)
			if vi == nil || vj == nil {
				continue
			}
			cmp := compareFieldValues(vi, vj)
			if cmp == 0 {
				continue
			}
			if ob.Descending {
				return cmp > 0
			}
			return cmp < 0
		}
		return false
	})
}

func findFieldValue(row []types.FieldValue, field string) interface{} {
	for _, fv := range row {
		if string(fv.FieldName) == field {
			return fv.Value
		}
	}
	return nil
}

func compareFieldValues(a, b interface{}) int {
	switch va := a.(type) {
	case uint64:
		vb := b.(uint64)
		if va < vb {
			return -1
		}
		if va > vb {
			return 1
		}
		return 0
	case int16:
		vb := a.(int16)
		vbb := b.(int16)
		if vb < vbb {
			return -1
		}
		if vb > vbb {
			return 1
		}
		return 0
	case int32:
		vb := b.(int32)
		if va < vb {
			return -1
		}
		if va > vb {
			return 1
		}
		return 0
	case byte:
		vb := b.(byte)
		if va < vb {
			return -1
		}
		if va > vb {
			return 1
		}
		return 0
	case []byte:
		vb := b.([]byte)
		return bytes.Compare(va, vb)
	}
	return 0
}

// ======================== 投影 ========================

func (q *Query) project(results [][]types.FieldValue) [][]types.FieldValue {
	colSet := make(map[string]bool)
	for _, c := range q.selectCols {
		colSet[c] = true
	}
	for i, row := range results {
		var filtered []types.FieldValue
		for _, fv := range row {
			if colSet[string(fv.FieldName)] {
				filtered = append(filtered, fv)
			}
		}
		results[i] = filtered
	}
	return results
}

// ======================== 流式聚合 ========================

// FindOne 流式查询，返回单行 (命中即停，不遍历全表)
func (q *Query) FindOne() ([]types.FieldValue, error) {
	var result []types.FieldValue
	found := false
	err := q.scan(func(row []types.FieldValue) bool {
		result = row
		found = true
		return false // 找到第一行即停止
	})
	if err != nil {
		return nil, err
	}
	if !found {
		return nil, page.ErrNotFound
	}
	return result, nil
}

// Count 流式计数 (不收集行，扫描过程中直接累加)
func (q *Query) Count() (int64, error) {
	var count int64
	err := q.scan(func(row []types.FieldValue) bool {
		count++
		return true
	})
	return count, err
}

// Sum 流式求和 (不收集行，扫描过程中直接累加)
func (q *Query) Sum(col string) (float64, error) {
	var sum float64
	err := q.scan(func(row []types.FieldValue) bool {
		for _, fv := range row {
			if string(fv.FieldName) == col {
				sum += toFloat64(fv.Value)
				break
			}
		}
		return true
	})
	return sum, err
}

// Max 流式求最大值 (不收集行，扫描过程中维护当前最大值)
func (q *Query) Max(col string) (interface{}, error) {
	var maxVal interface{}
	err := q.scan(func(row []types.FieldValue) bool {
		for _, fv := range row {
			if string(fv.FieldName) == col {
				if maxVal == nil || compareFieldValues(fv.Value, maxVal) > 0 {
					maxVal = fv.Value
				}
				break
			}
		}
		return true
	})
	return maxVal, err
}

// Min 流式求最小值 (不收集行，扫描过程中维护当前最小值)
func (q *Query) Min(col string) (interface{}, error) {
	var minVal interface{}
	err := q.scan(func(row []types.FieldValue) bool {
		for _, fv := range row {
			if string(fv.FieldName) == col {
				if minVal == nil || compareFieldValues(fv.Value, minVal) < 0 {
					minVal = fv.Value
				}
				break
			}
		}
		return true
	})
	return minVal, err
}

func toFloat64(v interface{}) float64 {
	switch val := v.(type) {
	case uint64:
		return float64(val)
	case int16:
		return float64(val)
	case int32:
		return float64(val)
	case byte:
		return float64(val)
	}
	return 0
}