// Package comp implements a byte code generator targeting the vm.
package comp

import (
	"fmt"
	"log"
	"os"
	"path"
	"reflect"
	"runtime"
	"strconv"

	"github.com/mvertes/parscan/lang"
	"github.com/mvertes/parscan/parser"
	"github.com/mvertes/parscan/scanner"
	"github.com/mvertes/parscan/symbol"
	"github.com/mvertes/parscan/vm"
)

// Compiler represents the state of a compiler.
type Compiler struct {
	*parser.Parser
	vm.Code            // produced code, to fill VM with
	Data    []vm.Value // produced data, will be at the bottom of VM stack
	Entry   int        // offset in Code to start execution from (skip function defintions)

	strings map[string]int // locations of strings in Data
}

// NewCompiler returns a new compiler state for a given scanner.
func NewCompiler(spec *lang.Spec) *Compiler {
	return &Compiler{
		Parser:  parser.NewParser(spec, true),
		Entry:   -1,
		strings: map[string]int{},
	}
}

// AddSym adds a new named value to the compiler symbol table, and returns its index in memory.
func (c *Compiler) AddSym(name string, value vm.Value) int {
	p := len(c.Data)
	c.Data = append(c.Data, value)
	c.Symbols.Add(p, name, value, symbol.Value, nil, false)
	return p
}

// Generate generates vm code and data from parsed tokens.
func (c *Compiler) Generate(tokens parser.Tokens) (err error) {
	log.Println("Codegen tokens:", tokens)
	fixList := parser.Tokens{}  // list of tokens to fix after all necessary information is gathered
	stack := []*symbol.Symbol{} // for symbolic evaluation, type checking, etc
	keyList := []string{}

	emit := func(t scanner.Token, op vm.Op, arg ...int) {
		_, file, line, _ := runtime.Caller(1)
		fmt.Fprintf(os.Stderr, "%s:%d: %v emit %v %v\n", path.Base(file), line, t, op, arg)
		c.Code = append(c.Code, vm.Instruction{Pos: vm.Pos(t.Pos), Op: op, Arg: arg})
	}
	push := func(s *symbol.Symbol) { stack = append(stack, s) }
	pop := func() *symbol.Symbol { l := len(stack) - 1; s := stack[l]; stack = stack[:l]; return s }
	top := func() *symbol.Symbol { return stack[len(stack)-1] }

	showStack := func() {
		_, file, line, _ := runtime.Caller(1)
		fmt.Fprintf(os.Stderr, "%s%d: showstack: %d\n", path.Base(file), line, len(stack))
		for i, s := range stack {
			fmt.Fprintf(os.Stderr, "%s:%d: stack[%d]: %v\n", path.Base(file), line, i, s)
		}
	}

	for i, t := range tokens {
		switch t.Tok {
		case lang.Int:
			n, err := strconv.Atoi(t.Str)
			if err != nil {
				return err
			}
			push(&symbol.Symbol{Kind: symbol.Const, Value: vm.ValueOf(n), Type: vm.TypeOf(0)})
			emit(t, vm.Push, n)

		case lang.String:
			s := t.Block()
			v := vm.Value{Type: vm.TypeOf(s), Value: reflect.ValueOf(s)}
			i, ok := c.strings[s]
			if !ok {
				i = len(c.Data)
				c.Data = append(c.Data, v)
				c.strings[s] = i
			}
			push(&symbol.Symbol{Kind: symbol.Const, Value: v})
			emit(t, vm.Dup, i)

		case lang.Add:
			push(&symbol.Symbol{Type: arithmeticOpType(pop(), pop())})
			emit(t, vm.Add)

		case lang.Mul:
			push(&symbol.Symbol{Type: arithmeticOpType(pop(), pop())})
			emit(t, vm.Mul)

		case lang.Sub:
			push(&symbol.Symbol{Type: arithmeticOpType(pop(), pop())})
			emit(t, vm.Sub)

		case lang.Minus:
			emit(t, vm.Push, 0)
			emit(t, vm.Sub)

		case lang.Not:
			emit(t, vm.Not)

		case lang.Plus:
			// Unary '+' is idempotent. Nothing to do.

		case lang.Addr:
			push(&symbol.Symbol{Type: vm.PointerTo(pop().Type)})
			emit(t, vm.Addr)

		case lang.Deref:
			push(&symbol.Symbol{Type: pop().Type.Elem()})
			emit(t, vm.Deref)

		case lang.Index:
			push(&symbol.Symbol{Type: pop().Type.Elem()})
			emit(t, vm.Index)

		case lang.Greater:
			push(&symbol.Symbol{Type: booleanOpType(pop(), pop())})
			emit(t, vm.Greater)

		case lang.Less:
			push(&symbol.Symbol{Type: booleanOpType(pop(), pop())})
			emit(t, vm.Lower)

		case lang.Call:
			s := pop()
			if s.Kind != symbol.Value {
				typ := s.Type
				// TODO: pop input types (careful with variadic function).
				for i := 0; i < typ.Rtype.NumOut(); i++ {
					push(&symbol.Symbol{Type: typ.Out(i)})
				}
				emit(t, vm.Call)
				break
			}
			push(s)
			fallthrough // A symValue must be called through callX.

		case lang.CallX:
			rtyp := pop().Value.Value.Type()
			// TODO: pop input types (careful with variadic function).
			for i := 0; i < rtyp.NumOut(); i++ {
				push(&symbol.Symbol{Type: &vm.Type{Rtype: rtyp.Out(i)}})
			}
			emit(t, vm.CallX, t.Beg)

		case lang.Colon:
			// Keyed element in a literal composite expression:
			// If the key is an ident (field name), then push it on the keystack,
			// to be computed at compile time in Composite handling.
			// Or generate instructions so key will be computed at runtime.
			if tokens[i-1].Tok == lang.Ident {
				keyList = append(keyList, tokens[i-1].Str)
			}

		case lang.Composite:
			d := top() // let the type symbol on the stack, may be required for assignment
			switch d.Type.Rtype.Kind() {
			case reflect.Struct:
				emit(t, vm.Fnew, d.Index)
				if len(keyList) == 0 {
					nf := d.Type.Rtype.NumField()
					for i := 0; i < nf; i++ {
						emit(t, vm.FieldSet, i)
					}
				} else {
					for _, fname := range keyList {
						i := d.Type.FieldIndex(fname)
						emit(t, vm.FieldSet, i...)
					}
					keyList = []string{}
				}
			case reflect.Slice:
				emit(t, vm.Fnew, d.Index)
				// if len(keyList) == 0 {
				// }
			default:
				return fmt.Errorf("composite kind not supported yet: %v", d.Type.Rtype.Kind())
			}

		case lang.Grow:
			emit(t, vm.Grow, t.Beg)

		case lang.Define:
			// TODO: support assignment to local, composite objects.
			st := tokens[i-1]
			l := len(c.Data)
			showStack()
			d := pop()
			typ := d.Type
			if typ == nil {
				typ = d.Value.Type
			}
			v := vm.NewValue(typ)
			c.Symbols.Add(l, st.Str, v, symbol.Var, typ, false)
			c.Data = append(c.Data, v)
			emit(t, vm.Assign, l)

		case lang.Assign:
			st := tokens[i-1]
			if st.Tok == lang.Period || st.Tok == lang.Index {
				emit(t, vm.Vassign)
				break
			}
			s, ok := c.Symbols[st.Str]
			if !ok {
				return fmt.Errorf("symbol not found: %s", st.Str)
			}
			d := pop()
			typ := d.Type
			if typ == nil {
				typ = d.Value.Type
			}
			if s.Type == nil {
				s.Type = typ
				s.Value = vm.NewValue(typ)
			}
			if s.Local {
				if !s.Used {
					emit(st, vm.New, s.Index, c.typeSym(s.Type).Index)
					s.Used = true
				}
				emit(st, vm.Fassign, s.Index)
				break
			}
			if s.Index == symbol.UnsetAddr {
				s.Index = len(c.Data)
				c.Data = append(c.Data, s.Value)
			}
			emit(st, vm.Assign, s.Index)

		case lang.Equal:
			push(&symbol.Symbol{Type: booleanOpType(pop(), pop())})
			emit(t, vm.Equal)

		case lang.EqualSet:
			push(&symbol.Symbol{Type: booleanOpType(pop(), pop())})
			emit(t, vm.EqualSet)

		case lang.Ident:
			if i < len(tokens)-1 {
				switch t1 := tokens[i+1]; t1.Tok {
				case lang.Define, lang.Assign, lang.Colon:
					continue
				}
			}
			s, ok := c.Symbols[t.Str]
			if !ok {
				return fmt.Errorf("symbol not found: %s", t.Str)
			}
			push(s)
			if s.Kind == symbol.Pkg {
				break
			}
			if s.Local {
				emit(t, vm.Fdup, s.Index)
			} else {
				if s.Index == symbol.UnsetAddr {
					s.Index = len(c.Data)
					c.Data = append(c.Data, s.Value)
				}
				if s.Kind != symbol.Type {
					emit(t, vm.Dup, s.Index)
				}
			}

		case lang.Label:
			lc := len(c.Code)
			s, ok := c.Symbols[t.Str]
			if ok {
				s.Value = vm.ValueOf(lc)
				if s.Kind == symbol.Func {
					// label is a function entry point, register its code address in data.
					s.Index = len(c.Data)
					c.Data = append(c.Data, s.Value)
				} else {
					c.Data[s.Index] = s.Value
				}
			} else {
				c.Symbols[t.Str] = &symbol.Symbol{Kind: symbol.Label, Value: vm.ValueOf(lc)}
			}

		case lang.JumpFalse:
			var i int
			if s, ok := c.Symbols[t.Str]; !ok {
				// t.Beg contains the position in code which needs to be fixed.
				t.Beg = len(c.Code)
				fixList = append(fixList, t)
			} else {
				i = int(s.Value.Int()) - len(c.Code)
			}
			emit(t, vm.JumpFalse, i)

		case lang.JumpSetFalse:
			var i int
			if s, ok := c.Symbols[t.Str]; !ok {
				// t.Beg contains the position in code which needs to be fixed.
				t.Beg = len(c.Code)
				fixList = append(fixList, t)
			} else {
				i = int(s.Value.Int()) - len(c.Code)
			}
			emit(t, vm.JumpSetFalse, i)

		case lang.JumpSetTrue:
			var i int
			if s, ok := c.Symbols[t.Str]; !ok {
				// t.Beg contains the position in code which needs to be fixed.
				t.Beg = len(c.Code)
				fixList = append(fixList, t)
			} else {
				i = int(s.Value.Int()) - len(c.Code)
			}
			emit(t, vm.JumpSetTrue, i)

		case lang.Goto:
			var i int
			if s, ok := c.Symbols[t.Str]; !ok {
				t.Beg = len(c.Code)
				fixList = append(fixList, t)
			} else {
				i = int(s.Value.Int()) - len(c.Code)
			}
			emit(t, vm.Jump, i)

		case lang.Period:
			s := pop()
			switch s.Kind {
			case symbol.Pkg:
				p, ok := parser.Packages[s.PkgPath]
				if !ok {
					return fmt.Errorf("package not found: %s", s.PkgPath)
				}
				v, ok := p[t.Str[1:]]
				if !ok {
					return fmt.Errorf("symbol not found in package %s: %s", s.PkgPath, t.Str[1:])
				}
				name := s.PkgPath + t.Str
				var l int
				sym, _, ok := c.Symbols.Get(name, "")
				if ok {
					l = sym.Index
				} else {
					l = len(c.Data)
					c.Data = append(c.Data, v)
					c.Symbols.Add(l, name, v, symbol.Value, v.Type, false)
					sym = c.Symbols[name]
				}
				push(sym)
				emit(t, vm.Dup, l)
			default:
				if f, ok := s.Type.Rtype.FieldByName(t.Str[1:]); ok {
					emit(t, vm.Field, f.Index...)
					break
				}
				return fmt.Errorf("field or method not found: %s", t.Str[1:])
			}

		case lang.Return:
			emit(t, vm.Return, t.Beg, t.End)

		default:
			return fmt.Errorf("generate: unsupported token %v", t)
		}
	}

	// Finally we fix unresolved labels for jump destinations.
	for _, t := range fixList {
		s, ok := c.Symbols[t.Str]
		if !ok {
			return fmt.Errorf("label not found: %q", t.Str)
		}
		c.Code[t.Beg].Arg[0] = int(s.Value.Int()) - t.Beg
	}
	return err
}
func arithmeticOpType(s1, _ *symbol.Symbol) *vm.Type { return symbol.Vtype(s1) }
func booleanOpType(_, _ *symbol.Symbol) *vm.Type     { return vm.TypeOf(true) }

// PrintCode pretty prints the generated code.
func (c *Compiler) PrintCode() {
	labels := map[int][]string{} // labels indexed by code location
	data := map[int]string{}     // data indexed by frame location

	for name, sym := range c.Symbols {
		if sym.Kind == symbol.Label || sym.Kind == symbol.Func {
			i := int(sym.Value.Int())
			labels[i] = append(labels[i], name)
		}
		if sym.Used {
			data[sym.Index] = name
		}
	}

	fmt.Fprintln(os.Stderr, "# Code:")
	for i, l := range c.Code {
		for _, label := range labels[i] {
			fmt.Fprintln(os.Stderr, label+":")
		}
		extra := ""
		switch l.Op {
		case vm.Jump, vm.JumpFalse, vm.JumpTrue, vm.JumpSetFalse, vm.JumpSetTrue, vm.Calli:
			if d, ok := labels[i+l.Arg[0]]; ok {
				extra = "// " + d[0]
			}
		case vm.Dup, vm.Assign, vm.Fdup, vm.Fassign:
			if d, ok := data[l.Arg[0]]; ok {
				extra = "// " + d
			}
		}
		fmt.Fprintf(os.Stderr, "%4d %v %v\n", i, l, extra)
	}

	for _, label := range labels[len(c.Code)] {
		fmt.Fprintln(os.Stderr, label+":")
	}
	fmt.Fprintln(os.Stderr, "# End code")
}

type entry struct {
	name string
	*symbol.Symbol
}

func (e entry) String() string { return fmt.Sprintf("name: %s, sym: %v", e.name, e.Symbol) }

// PrintData pretty prints the generated global data symbols in compiler.
func (c *Compiler) PrintData() {
	dict := c.symbolsByIndex()

	fmt.Fprintln(os.Stderr, "# Data:")
	for i, d := range c.Data {
		fmt.Fprintf(os.Stderr, "%4d %T %v %v\n", i, d.Interface(), d.Value, dict[i])
	}
}

func (c *Compiler) symbolsByIndex() map[int]entry {
	dict := map[int]entry{}
	for name, sym := range c.Symbols {
		if sym.Index == symbol.UnsetAddr {
			continue
		}
		dict[sym.Index] = entry{name, sym}
	}
	return dict
}

// Dump represents the state of a data dump.
type Dump struct {
	Values []*DumpValue
}

// DumpValue is a value of a dump state.
type DumpValue struct {
	Index int
	Name  string
	Kind  int
	Type  string
	Value any
}

// Dump creates a snapshot of the execution state of global variables.
// This method is specifically implemented in the Compiler to minimize the coupling between
// the dump format and other components. By situating the dump logic in the Compiler,
// it relies solely on the program being executed and the indexing algorithm used for ordering variables
// (currently, this is an integer that corresponds to the order of variables in the program).
// This design choice allows the Virtual Machine (VM) to evolve its memory management strategies
// without compromising backward compatibility with dumps generated by previous versions.
func (c *Compiler) Dump() *Dump {
	dict := c.symbolsByIndex()
	dv := make([]*DumpValue, len(c.Data))
	for i, d := range c.Data {
		e := dict[i]
		dv[i] = &DumpValue{
			Index: e.Index,
			Name:  e.name,
			Kind:  int(e.Kind),
			Type:  e.Type.Name,
			Value: d.Interface(),
		}
	}
	return &Dump{Values: dv}
}

// ApplyDump sets previously saved dump, restoring the state of global variables.
func (c *Compiler) ApplyDump(d *Dump) error {
	dict := c.symbolsByIndex()
	for _, dv := range d.Values {
		// do all the checks to be sure we are applying the correct values
		e, ok := dict[dv.Index]
		if !ok {
			return fmt.Errorf("entry not found on index %d", dv.Index)
		}

		if dv.Name != e.name ||
			dv.Type != e.Type.Name ||
			dv.Kind != int(e.Kind) {
			return fmt.Errorf("entry with index %d does not match with provided entry. "+
				"dumpValue: %s, %s, %d. memoryValue: %s, %s, %d",
				dv.Index,
				dv.Name, dv.Type, dv.Kind,
				e.name, e.Type, e.Kind)
		}

		if dv.Index >= len(c.Data) {
			return fmt.Errorf("index (%d) bigger than memory (%d)", dv.Index, len(c.Data))
		}

		if !c.Data[dv.Index].CanSet() {
			return fmt.Errorf("value %v cannot be set", dv.Value)
		}

		c.Data[dv.Index].Set(reflect.ValueOf(dv.Value))
	}
	return nil
}

func (c *Compiler) typeSym(t *vm.Type) *symbol.Symbol {
	tsym, ok := c.Symbols[t.Rtype.String()]
	if !ok {
		tsym = &symbol.Symbol{Index: symbol.UnsetAddr, Kind: symbol.Type, Type: t}
		c.Symbols[t.Rtype.String()] = tsym
	}
	if tsym.Index == symbol.UnsetAddr {
		tsym.Index = len(c.Data)
		c.Data = append(c.Data, vm.NewValue(t))
	}
	return tsym
}