package grammar

import (
	"fmt"
)

type followEntry struct {
	symbols map[symbol]struct{}
	eof     bool
}

func newFollowEntry() *followEntry {
	return &followEntry{
		symbols: map[symbol]struct{}{},
		eof:     false,
	}
}

func (e *followEntry) add(sym symbol) bool {
	if _, ok := e.symbols[sym]; ok {
		return false
	}
	e.symbols[sym] = struct{}{}
	return true
}

func (e *followEntry) addEOF() bool {
	if !e.eof {
		e.eof = true
		return true
	}
	return false
}

func (e *followEntry) merge(fst *firstEntry, flw *followEntry) bool {
	changed := false

	if fst != nil {
		for sym := range fst.symbols {
			added := e.add(sym)
			if added {
				changed = true
			}
		}
	}

	if flw != nil {
		for sym := range flw.symbols {
			added := e.add(sym)
			if added {
				changed = true
			}
		}
		if flw.eof {
			added := e.addEOF()
			if added {
				changed = true
			}
		}
	}

	return changed
}

type followSet struct {
	set map[symbol]*followEntry
}

func newFollow(prods *productionSet) *followSet {
	flw := &followSet{
		set: map[symbol]*followEntry{},
	}
	for _, prod := range prods.getAllProductions() {
		if _, ok := flw.set[prod.lhs]; ok {
			continue
		}
		flw.set[prod.lhs] = newFollowEntry()
	}
	return flw
}

func (flw *followSet) find(sym symbol) (*followEntry, error) {
	e, ok := flw.set[sym]
	if !ok {
		return nil, fmt.Errorf("an entry of FOLLOW was not found; symbol: %s", sym)
	}
	return e, nil
}

func genFollowSet(prods *productionSet, first *firstSet) (*followSet, error) {
	ntsyms := map[symbol]struct{}{}
	for _, prod := range prods.getAllProductions() {
		if _, ok := ntsyms[prod.lhs]; ok {
			continue
		}
		ntsyms[prod.lhs] = struct{}{}
	}

	follow := newFollow(prods)
	for {
		more := false
		for ntsym := range ntsyms {
			e, err := follow.find(ntsym)
			if err != nil {
				return nil, err
			}
			if ntsym.isStart() {
				changed := e.addEOF()
				if changed {
					more = true
				}
			}
			for _, prod := range prods.getAllProductions() {
				for i, sym := range prod.rhs {
					if sym != ntsym {
						continue
					}
					fst, err := first.find(prod, i+1)
					if err != nil {
						return nil, err
					}
					changed := e.merge(fst, nil)
					if changed {
						more = true
					}
					if fst.empty {
						flw, err := follow.find(prod.lhs)
						if err != nil {
							return nil, err
						}
						changed := e.merge(nil, flw)
						if changed {
							more = true
						}
					}
				}
			}
		}
		if !more {
			break
		}
	}

	return follow, nil
}