package subex

import (
	"main/walk"
)

type Transducer struct {
	storeSize    NextSlotIds
	initialState SubexState
}

// Where slots are stored
type Store struct {
	values [][]walk.Value
	runes  [][]rune
}

// Return a new store with all the data from this one
func (store Store) clone() Store {
	newStore := Store{
		values: make([][]walk.Value, len(store.values)),
		runes:  make([][]rune, len(store.runes)),
	}
	copy(newStore.values, store.values)
	copy(newStore.runes, store.runes)
	return newStore
}

// Return a copy of this store but with an additional slot set
func (store Store) withValue(key int, value []walk.Value) Store {
	newStore := store.clone()
	newStore.values[key] = value
	return newStore
}

func (store Store) withRunes(key int, runes []rune) Store {
	newStore := store.clone()
	newStore.runes[key] = runes
	return newStore
}

type SlotId struct {
	id  int
	typ Type
}

type NextSlotIds struct {
	values int
	runes  int
}

type SlotMap struct {
	next NextSlotIds
	ids  map[rune]SlotId
}

func (m *SlotMap) getId(slot rune) int {
	id, exists := m.ids[slot]
	if exists {
		if id.typ != ValueType {
			panic("Slot with wrong type used")
		}
		return id.id
	}
	id.id = m.next.values
	id.typ = ValueType
	m.next.values++
	m.ids[slot] = id
	return id.id
}
func (m *SlotMap) getRuneId(slot rune) int {
	id, exists := m.ids[slot]
	if exists {
		if id.typ != RuneType {
			panic("Slot with wrong type used")
		}
		return id.id
	}
	id.id = m.next.runes
	id.typ = RuneType
	m.next.runes++
	m.ids[slot] = id
	return id.id
}

// Compile the SubexAST into a transducer SubexState that can be run
func CompileTransducer(transducerAst SubexAST) Transducer {
	slotMap := SlotMap{
		next: NextSlotIds{
			values: 0,
			runes:  0,
		},
		ids: make(map[rune]SlotId),
	}
	initial := transducerAst.compileWith(&SubexNoneState{}, &slotMap, ValueType, ValueType)
	return Transducer{
		storeSize:    slotMap.next,
		initialState: initial,
	}
}

// An immutable stack for outputting to
type OutputStack struct {
	head walk.OutputList
	tail *OutputStack
}

func (stack OutputStack) pop() ([]walk.Value, OutputStack) {
	return stack.head.(walk.OutputValueList), *stack.tail
}
func (stack OutputStack) push(atoms []walk.Value) OutputStack {
	return OutputStack{
		head: walk.OutputValueList(atoms),
		tail: &stack,
	}
}
func (stack OutputStack) replace(atoms []walk.Value) OutputStack {
	return OutputStack{
		head: walk.OutputValueList(atoms),
		tail: stack.tail,
	}
}
func (stack OutputStack) peek() []walk.Value {
	return stack.head.(walk.OutputValueList)
}

func topAppend(outputStack OutputStack, values []walk.Value) OutputStack {
	head := outputStack.peek()
	head = append([]walk.Value{}, head...)
	head = append(head, values...)
	return outputStack.replace(head)
}

func topAppendRune(outputStack OutputStack, runes []rune) OutputStack {
	head := outputStack.head.(walk.OutputRuneList)
	head = append([]rune{}, head...)
	head = append(head, runes...)
	return OutputStack{
		head: head,
		tail: outputStack.tail,
	}
}

// Additional state that goes along with a subex state in an execution branch
type auxiliaryState struct {
	// Content of slots in this branch
	store Store
	// The output stack. At the end of the program, whatever is on top of this will be output
	// States may push or pop to the stack as they wish, creating sort of a call stack that allows states to capture the output of other states
	outputStack OutputStack
	// How deeply nested the current execution is inside of the overall value
	// i.e. starts at zero, is incremented to one when entering an array
	nesting int
}

func (aux auxiliaryState) cloneStore() auxiliaryState {
	aux.store = aux.store.clone()
	return aux
}

func (aux auxiliaryState) withValue(slot int, value []walk.Value) auxiliaryState {
	aux.store = aux.store.withValue(slot, value)
	return aux
}

func (aux auxiliaryState) pushOutput(data []walk.Value) auxiliaryState {
	aux.outputStack = aux.outputStack.push(data)
	return aux
}

func (aux auxiliaryState) pushOutputRunes(runes []rune) auxiliaryState {
	tail := aux.outputStack
	aux.outputStack = OutputStack{
		head: walk.OutputRuneList(runes),
		tail: &tail,
	}
	return aux
}

func (aux auxiliaryState) popDiscardOutput() auxiliaryState {
	aux.outputStack = *aux.outputStack.tail
	return aux
}

func (aux auxiliaryState) popOutput() ([]walk.Value, auxiliaryState) {
	data, output := aux.outputStack.pop()
	aux.outputStack = output
	return data, aux
}

func (aux auxiliaryState) popOutputRunes() ([]rune, auxiliaryState) {
	runes := aux.outputStack.head.(walk.OutputRuneList)
	aux.outputStack = *aux.outputStack.tail
	return runes, aux
}

func (aux auxiliaryState) topAppend(values []walk.Value) auxiliaryState {
	aux.outputStack = topAppend(aux.outputStack, values)
	return aux
}

func (aux auxiliaryState) topAppendRune(runes []rune) auxiliaryState {
	aux.outputStack = topAppendRune(aux.outputStack, runes)
	return aux
}

func (aux auxiliaryState) incNest() auxiliaryState {
	aux.nesting++
	return aux
}

func (aux auxiliaryState) decNest() auxiliaryState {
	aux.nesting--
	return aux
}

type SubexBranch struct {
	state SubexState
	aux   auxiliaryState
}

// One branch of subex execution
type SubexEatBranch struct {
	// State in this branch
	state SubexEatState
	// Axiliary state
	aux auxiliaryState
}

// Read a single character and return all the branches resulting from this branch consuming it
func (pair SubexEatBranch) eat(edible walk.Edible) []SubexBranch {
	return pair.state.eat(pair.aux, edible)
}
func (pair SubexEatBranch) accepting() []OutputStack {
	return pair.state.accepting(pair.aux)
}

func equalStates(left SubexEatBranch, right SubexEatBranch) bool {
	// Only care about if they are the same pointer
	return left.state == right.state && left.aux.nesting == right.aux.nesting
}

// If two branches have the same state, only the first has a chance of being successful
// This function removes all of the pointless execution branches to save execution time
func pruneStates(states []SubexEatBranch) []SubexEatBranch {
	uniqueStates := 0
outer:
	for _, state := range states {
		for i := 0; i < uniqueStates; i++ {
			if equalStates(state, states[i]) {
				continue outer
			}
		}
		states[uniqueStates] = state
		uniqueStates++
	}
	return states[:uniqueStates]
}

func addStates(curStates []SubexEatBranch, newStates []SubexBranch) []SubexEatBranch {
	for _, state := range newStates {
		switch s := state.state.(type) {
		case SubexEpsilonState:
			curStates = addStates(curStates, s.epsilon(state.aux))
		case SubexEatState:
			curStates = append(curStates, SubexEatBranch{
				state: s,
				aux:   state.aux,
			})
		}
	}
	return curStates
}

func processInput(states []SubexEatBranch, input walk.Edible, nesting int) []SubexEatBranch {
	newStates := make([]SubexEatBranch, 0, 2)

	for _, state := range states {
		// TODO: What if nesting is changed by an epsilon state?
		if state.aux.nesting == nesting {
			newStates = addStates(newStates, state.eat(input))
		} else if state.aux.nesting < nesting {
			newStates = append(newStates, state)
		}
	}

	switch input := input.(type) {
	case walk.StringValue:
		for _, r := range input {
			newStates = processInput(newStates, walk.RuneEdible(r), nesting+1)
		}
		newStates = processInput(newStates, walk.StringEnd, nesting+1)
	case walk.ArrayValue:
		for _, el := range input {
			newStates = processInput(newStates, walk.NumberValue(el.Index), nesting+1)
			newStates = processInput(newStates, el.Value, nesting+1)
		}
		newStates = processInput(newStates, walk.ArrayEnd, nesting+1)
	case walk.MapValue:
		for _, el := range input {
			newStates = processInput(newStates, walk.StringValue(el.Key), nesting+1)
			newStates = processInput(newStates, el.Value, nesting+1)
		}
		newStates = processInput(newStates, walk.MapEnd, nesting+1)
	}

	newStates = pruneStates(newStates)

	return newStates
}

// Run the subex transducer
func RunTransducer(transducer Transducer, input []walk.Value) (output []walk.Value, err bool) {
	states := addStates(nil, []SubexBranch{{
		state: transducer.initialState,
		aux: auxiliaryState{
			outputStack: OutputStack{
				head: walk.OutputValueList(nil),
				tail: nil,
			},
			store: Store{
				values: make([][]walk.Value, transducer.storeSize.values),
				runes:  make([][]rune, transducer.storeSize.runes),
			},
			nesting: 0,
		},
	}})

	for _, value := range input {
		if len(states) == 0 {
			break
		}

		states = processInput(states, value, 0)
	}

	for _, state := range states {
		if state.aux.nesting > 0 {
			continue
		}
		acceptingStacks := state.accepting()
		for _, stack := range acceptingStacks {
			return stack.head.(walk.OutputValueList), false
		}
	}
	return nil, true
}