// Copyright 2021 PingCAP, 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 utils

import (
	"context"
	"errors"

	. "github.com/pingcap/endless/pkg/matcher"
)

type Iterator interface {
	Valid() bool
	Value() interface{}
	Next() error
	Close()
}

type SimpleIterator struct {
	items []interface{}
	cur   int
}

func NewSimpleIterator(items []interface{}) *SimpleIterator {
	return &SimpleIterator{
		items: items,
		cur:   0,
	}
}

func (i *SimpleIterator) Valid() bool {
	return i.cur >= 0 && i.cur < len(i.items)
}

func (i *SimpleIterator) Value() interface{} {
	if i.Valid() {
		return i.items[i.cur]
	}
	return nil
}

func (i *SimpleIterator) Next() error {
	if !i.Valid() {
		return errors.New("iter is invalid")
	}
	i.cur++
	return nil
}

func (i *SimpleIterator) Close() {
	i.cur = -1
}

type Forker interface {
	DoFork(ctx context.Context) (Iterator, error)
}

type FlatForker struct {
	items []interface{}
}

func NewFlatForker(items []interface{}) *FlatForker {
	return &FlatForker{items: items}
}

func (f *FlatForker) DoFork(_ context.Context) (Iterator, error) {
	return NewSimpleIterator(f.items), nil
}

type GenerationForker struct {
	fn func(ctx *GenerateContext) (interface{}, error)
}

func NewGenerationForker(fn func(ctx *GenerateContext) (interface{}, error)) *GenerationForker {
	return &GenerationForker{fn: fn}
}

func (f *GenerationForker) DoFork(ctx context.Context) (Iterator, error) {
	return newGenerator(ctx, f.fn)
}

type generator struct {
	ctx       context.Context
	cancelCtx func()
	fn        func(ctx *GenerateContext) (interface{}, error)
	valid     bool
	value     interface{}
	stack     []Iterator
}

func newGenerator(ctx context.Context, fn func(ctx *GenerateContext) (interface{}, error)) (*generator, error) {
	newCtx, cancel := context.WithCancel(ctx)
	gen := &generator{
		ctx:       newCtx,
		cancelCtx: cancel,
		fn:        fn,
		value:     nil,
		valid:     true,
	}
	if err := gen.next(); err != nil {
		return nil, err
	}
	return gen, nil
}

func (g *generator) Valid() bool {
	return g.valid
}

func (g *generator) Value() interface{} {
	if g.valid {
		return g.value
	}
	return nil
}

func (g *generator) Next() error {
	return g.next()
}

func (g *generator) Close() {
	if g.valid {
		g.cancelCtx()
		g.valid = false
		g.value = nil
		g.cancelCtx = nil
	}
}

func (g *generator) next() error {
	if !g.valid {
		return errors.New("not valid")
	}

	if g.stack == nil {
		g.stack = make([]Iterator, 0)
	} else {
		for len(g.stack) > 0 {
			last := len(g.stack) - 1
			iter := g.stack[last]
			if err := iter.Next(); err != nil {
				g.Close()
				return err
			}

			if iter.Valid() {
				break
			}
			g.stack = g.stack[0:last]
		}

		if len(g.stack) == 0 {
			g.Close()
			return nil
		}
	}

	fixedValues := make([]interface{}, 0, len(g.stack))
	for _, item := range g.stack {
		fixedValues = append(fixedValues, item.Value())
	}

	v, err := g.fn(&GenerateContext{generator: g, fixedValues: fixedValues, nextValuePos: 0})
	if err != nil {
		g.Close()
		return err
	}

	g.value = v
	return nil
}

type GenerateContext struct {
	generator    *generator
	fixedValues  []interface{}
	nextValuePos int
}

func (ctx *GenerateContext) Pick(from interface{}) (interface{}, error) {
	if pos := ctx.nextValuePos; pos < len(ctx.fixedValues) {
		ctx.nextValuePos++
		return ctx.fixedValues[pos], nil
	}

	forker, ok := from.(Forker)
	if !ok {
		forker = NewFlatForker([]interface{}{from})
	}

	iter, err := forker.DoFork(ctx.generator.ctx)
	if err != nil {
		return nil, err
	}

	if !iter.Valid() {
		return nil, errors.New("forker is empty")
	}
	ctx.generator.stack = append(ctx.generator.stack, iter)
	return iter.Value(), nil
}

func (ctx *GenerateContext) PickEnum(value ...interface{}) (interface{}, error) {
	return ctx.Pick(NewFlatForker(value))
}

func (ctx *GenerateContext) PickIntEnum(value ...int) (int, error) {
	values := make([]interface{}, 0, len(value))
	for _, v := range value {
		values = append(values, v)
	}

	v, err := ctx.PickEnum(values...)
	if err != nil {
		return 0, err
	}
	return v.(int), nil
}

func (ctx *GenerateContext) PickStringEnum(value ...string) (string, error) {
	values := make([]interface{}, 0, len(value))
	for _, v := range value {
		values = append(values, v)
	}

	v, err := ctx.PickEnum(values...)
	if err != nil {
		return "", err
	}
	return v.(string), nil
}

func (ctx *GenerateContext) PickBool() (bool, error) {
	v, err := ctx.Pick(NewFlatForker([]interface{}{false, true}))
	if err != nil {
		return false, err
	}
	return v.(bool), nil
}

func ForkRun(fn func(ctx *GenerateContext)) {
	iter := Try(NewGenerationForker(func(ctx *GenerateContext) (interface{}, error) {
		fn(ctx)
		return nil, nil
	}).DoFork(context.TODO())).(Iterator)

	for iter.Valid() {
		Try(iter.Next())
	}
}