// Copyright 2024 Dolthub, Inc.
//
// Licensed 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 val

import (
	"context"
)

// ExtendedTupleComparator is a comparator that properly handles extended types.
type ExtendedTupleComparator struct {
	innerCmp TupleComparator
	handlers []TupleTypeHandler
	vs       ValueStore
}

// TODO: compare performance of rolling this logic into the DefaultTupleComparator (nil check or generic handlers that call compare)
var _ TupleComparator = &ExtendedTupleComparator{}

// Compare implements the TupleComparator interface.
func (c *ExtendedTupleComparator) Compare(ctx context.Context, left, right Tuple, desc *TupleDesc) (cmp int, err error) {
	fast := desc.GetFixedAccess()
	off := len(fast)
	var start, stop ByteSize
	for i := 0; i < off; i++ {
		stop = fast[i]
		cmp, err = c.CompareValues(ctx, i, left[start:stop], right[start:stop], desc.Types[i])
		if err != nil {
			return 0, err
		}
		if cmp != 0 {
			return cmp, nil
		}
		start = stop
	}

	for i, typ := range desc.Types[off:] {
		j := i + off
		cmp, err = c.CompareValues(ctx, j, left.GetField(j), right.GetField(j), typ)
		if err != nil {
			return 0, err
		}
		if cmp != 0 {
			return cmp, nil
		}
	}
	return
}

// CompareValues implements the TupleComparator interface.
func (c *ExtendedTupleComparator) CompareValues(ctx context.Context, index int, left, right []byte, typ Type) (int, error) {
	switch typ.Enc {
	case ExtendedEnc, ExtendedAddrEnc, ExtendedAdaptiveEnc:
		return c.handlers[index].SerializedCompare(ctx, left, right)
	default:
		return compare(ctx, typ, left, right, c.vs)
	}
}

// Prefix implements the TupleComparator interface.
func (c *ExtendedTupleComparator) Prefix(n int) TupleComparator {
	return &ExtendedTupleComparator{innerCmp: c.innerCmp.Prefix(n), handlers: c.handlers[:n], vs: c.vs}
}

// Suffix implements the TupleComparator interface.
func (c *ExtendedTupleComparator) Suffix(n int) TupleComparator {
	return &ExtendedTupleComparator{innerCmp: c.innerCmp.Suffix(n), handlers: c.handlers[n:], vs: c.vs}
}

// Validated implements the TupleComparator interface.
func (c *ExtendedTupleComparator) Validated(types []Type) TupleComparator {
	// If our inner comparator is an ExtendedTupleComparator, then we should use its inner comparator to reduce redundancy, as well as its ValueStore if set.
	var innerCmp TupleComparator
	vs := c.vs
	if extendedInner, ok := c.innerCmp.(*ExtendedTupleComparator); ok {
		innerCmp = extendedInner.innerCmp.Validated(types)
		if vs == nil {
			vs = extendedInner.vs
		}
	} else {
		innerCmp = c.innerCmp.Validated(types)
	}
	if len(c.handlers) == 0 {
		c.handlers = make([]TupleTypeHandler, len(types))
	} else if len(c.handlers) != len(types) {
		panic("invalid handler count compared to types")
	}
	hasHandler := false
	for i, handler := range c.handlers {
		switch types[i].Enc {
		case ExtendedEnc, ExtendedAddrEnc, ExtendedAdaptiveEnc:
			if handler == nil {
				panic("extended encoding requires a handler")
			} else {
				hasHandler = true
			}
		}
	}
	if !hasHandler {
		return innerCmp
	}
	return &ExtendedTupleComparator{innerCmp: innerCmp, handlers: c.handlers, vs: vs}
}

// WithValueStore implements the TupleComparator interface.
func (c *ExtendedTupleComparator) WithValueStore(vs ValueStore) TupleComparator {
	return &ExtendedTupleComparator{
		innerCmp: c.innerCmp.WithValueStore(vs),
		handlers: c.handlers,
		vs:       vs,
	}
}