package schema

import (
	"github.com/go-gota/gota/dataframe"
	"github.com/go-gota/gota/series"
	"github.com/pingcap-inc/trino-thrift-connectors/gen-go/trino"
	"github.com/samber/lo"
)

type SplitInfo struct {
	SchemaTableName  *trino.TrinoThriftSchemaTableName
	OutputConstraint *trino.TrinoThriftTupleDomain
	IndexColumnNames []string
	Keys             *trino.TrinoThriftPageResult_
	PartInfo         any
}

func (si *SplitInfo) Filter(df dataframe.DataFrame) []int {
	if df.Nrow() == 0 {
		return []int{}
	}

	indexes := []int{}
	for idx := 0; idx < df.Nrow(); idx++ {
		indexes = append(indexes, idx)
	}
	df = df.Mutate(series.Ints(indexes))
	columns := df.Names()

	if len(si.IndexColumnNames) > 0 {
		f := si.toFrame(si.IndexColumnNames, si.Keys.ColumnBlocks)
		df = df.InnerJoin(f, si.IndexColumnNames...)
	}

	indexes, _ = df.Col(columns[len(columns)-1]).Int()
	for column, domain := range si.OutputConstraint.Domains {
		if domain.ValueSet == nil || df.Col(column).Err != nil {
			continue
		}

		if domain.ValueSet.EquatableValueSet != nil && domain.ValueSet.EquatableValueSet.Inclusive {
			f := si.toFrame([]string{column}, domain.ValueSet.EquatableValueSet.Values)
			df = df.Filter(dataframe.F{0, column, series.In, f.Col(column)})
			xs, _ := df.Col(columns[len(columns)-1]).Int()
			indexes = lo.Intersect(indexes, xs)
		}

		if domain.ValueSet.RangeValueSet != nil {
			domain_indexes := []int{}
			for _, r := range domain.ValueSet.RangeValueSet.Ranges {
				f := df.Copy()
				if r.Low.Value != nil {
					filters := []dataframe.F{}
					lowValue := si.toSeries(r.Low.Value).Val(0)
					filters = append(filters, dataframe.F{0, column, series.Greater, lowValue})
					if r.Low.Bound == trino.TrinoThriftBound_EXACTLY {
						filters = append(filters, dataframe.F{0, column, series.Eq, lowValue})
					}
					f = f.Filter(filters...)
				}
				if r.High.Value != nil {
					filters := []dataframe.F{}
					highValue := si.toSeries(r.High.Value).Val(0)
					filters = append(filters, dataframe.F{0, column, series.Less, highValue})
					if r.High.Bound == trino.TrinoThriftBound_EXACTLY {
						filters = append(filters, dataframe.F{0, column, series.Eq, highValue})
					}
					f = f.Filter(filters...)
				}
				xs, _ := f.Col(columns[len(columns)-1]).Int()
				domain_indexes = append(domain_indexes, xs...)
			}
			indexes = lo.Intersect(indexes, domain_indexes)
		}
	}

	return indexes
}

func (si *SplitInfo) toSeries(block *trino.TrinoThriftBlock) series.Series {
	if block.IntegerData != nil {
		return series.Ints(lo.Map(block.IntegerData.Ints, func(x int32, i int) int { return int(x) }))
	} else if block.BigintData != nil {
		return series.Ints(lo.Map(block.BigintData.Longs, func(x int64, i int) int { return int(x) }))
	} else if block.DoubleData != nil {
		return series.Floats(block.DoubleData.Doubles)
	} else if block.VarcharData != nil {
		strings := []string{}
		buf := block.VarcharData.Bytes
		for _, size := range block.VarcharData.Sizes {
			strings = append(strings, string(buf[:size]))
			buf = buf[size:]
		}
		return series.Strings(strings)
	} else if block.BooleanData != nil {
		return series.Bools(block.BooleanData.Booleans)
	} else if block.DateData != nil {
		return series.Ints(lo.Map(block.DateData.Dates, func(x int32, i int) int { return int(x) }))
	} else if block.TimestampData != nil {
		return series.Ints(lo.Map(block.TimestampData.Timestamps, func(x int64, i int) int { return int(x) }))
	} else if block.JsonData != nil {
		strings := []string{}
		buf := block.JsonData.Bytes
		for _, size := range block.JsonData.Sizes {
			strings = append(strings, string(buf[:size]))
			buf = buf[size:]
		}
		return series.Strings(strings)
	} else if block.HyperLogLogData != nil {
		strings := []string{}
		buf := block.HyperLogLogData.Bytes
		for _, size := range block.HyperLogLogData.Sizes {
			strings = append(strings, string(buf[:size]))
			buf = buf[size:]
		}
		return series.Strings(strings)
	} else if block.BigintArrayData != nil {
		return series.Ints(lo.Map(block.BigintArrayData.Values.Longs, func(x int64, i int) int { return int(x) }))
	}

	panic("unreachable")
}

func (si *SplitInfo) toFrame(columns []string, blocks []*trino.TrinoThriftBlock) dataframe.DataFrame {
	ss := []series.Series{}

	for idx, block := range blocks {
		s := si.toSeries(block)
		s.Name = columns[idx]
		ss = append(ss, s)
	}

	return dataframe.New(ss...)
}