package driver

import (
	_ "embed"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"os"
	"strings"
	"text/template"

	"github.com/nihei9/maleeni/spec"
)

//go:embed lexer.go
var lexerCoreSrc string

func GenLexer(clspec *spec.CompiledLexSpec, pkgName string) error {
	var lexerSrc string
	{
		fset := token.NewFileSet()
		f, err := parser.ParseFile(fset, "lexer.go", lexerCoreSrc, parser.ParseComments)
		if err != nil {
			return err
		}

		var b strings.Builder
		err = format.Node(&b, fset, f)
		if err != nil {
			return err
		}

		lexerSrc = b.String()
	}

	var specSrc string
	{
		t, err := template.New("").Funcs(genTemplateFuncs(clspec)).Parse(lexSpecTemplate)
		if err != nil {
			return err
		}

		var b strings.Builder
		err = t.Execute(&b, map[string]interface{}{
			"initialModeID":    clspec.InitialModeID,
			"modeIDNil":        spec.LexModeIDNil,
			"modeKindIDNil":    spec.LexModeKindIDNil,
			"stateIDNil":       spec.StateIDNil,
			"compressionLevel": clspec.CompressionLevel,
		})
		if err != nil {
			return err
		}

		specSrc = b.String()
	}

	var src string
	{
		tmpl := `// Code generated by maleeni-go. DO NOT EDIT.
{{ .lexerSrc }}

{{ .specSrc }}
`

		t, err := template.New("").Parse(tmpl)
		if err != nil {
			return err
		}

		var b strings.Builder
		err = t.Execute(&b, map[string]string{
			"lexerSrc": lexerSrc,
			"specSrc":  specSrc,
		})
		if err != nil {
			return err
		}

		src = b.String()
	}

	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, "", src, parser.ParseComments)
	if err != nil {
		return err
	}

	f.Name = ast.NewIdent(pkgName)

	return format.Node(os.Stdout, fset, f)
}

const lexSpecTemplate = `
type lexSpec struct {
	pop           [][]bool
	push          [][]ModeID
	modeNames     []string
	initialStates []StateID
	acceptances   [][]ModeKindID
	kindIDs       [][]KindID
	kindNames     []string
	initialModeID ModeID
	modeIDNil     ModeID
	modeKindIDNil ModeKindID
	stateIDNil    StateID

	rowNums           [][]int
	rowDisplacements  [][]int
	bounds            [][]int
	entries           [][]StateID
	originalColCounts []int
}

func NewLexSpec() *lexSpec {
	return &lexSpec{
		pop: {{ genPopTable }},
		push: {{ genPushTable }},
		modeNames: {{ genModeNameTable }},
		initialStates: {{ genInitialStateTable }},
		acceptances: {{ genAcceptTable }},
		kindIDs: {{ genKindIDTable }},
		kindNames: {{ genKindNameTable }},
		initialModeID: {{ .initialModeID }},
		modeIDNil: {{ .modeIDNil }},
		modeKindIDNil: {{ .modeKindIDNil }},
		stateIDNil: {{ .stateIDNil }},

		rowNums: {{ genRowNums }},
		rowDisplacements: {{ genRowDisplacements }},
		bounds: {{ genBounds }},
		entries: {{ genEntries }},
		originalColCounts: {{ genOriginalColCounts }},
	}
}

func (s *lexSpec) InitialMode() ModeID {
	return s.initialModeID
}

func (s *lexSpec) Pop(mode ModeID, modeKind ModeKindID) bool {
	return s.pop[mode][modeKind]
}

func (s *lexSpec) Push(mode ModeID, modeKind ModeKindID) (ModeID, bool) {
	id := s.push[mode][modeKind]
	return id, id != s.modeIDNil
}

func (s *lexSpec) ModeName(mode ModeID) string {
	return s.modeNames[mode]
}

func (s *lexSpec) InitialState(mode ModeID) StateID {
	return s.initialStates[mode]
}

func (s *lexSpec) NextState(mode ModeID, state StateID, v int) (StateID, bool) {
{{ if eq .compressionLevel 2 -}}
	rowNum := s.rowNums[mode][state]
	d := s.rowDisplacements[mode][rowNum]
	if s.bounds[mode][d+v] != rowNum {
		return s.stateIDNil, false
	}
	return s.entries[mode][d+v], true
{{ else if eq .compressionLevel 1 -}}
	rowNum := s.rowNums[mode][state]
	colCount := s.originalColCounts[mode]
	next := s.entries[mode][rowNum*colCount+v]
	if next == s.stateIDNil {
		return s.stateIDNil, false
	}
	return next, true
{{ else -}}
	colCount := s.originalColCounts[mode]
	next := s.entries[mode][int(state)*colCount+v]
	if next == s.stateIDNil {
		return s.stateIDNil, false
	}
	return next, true
{{ end -}}
}

func (s *lexSpec) Accept(mode ModeID, state StateID) (ModeKindID, bool) {
	id := s.acceptances[mode][state]
	return id, id != s.modeKindIDNil
}

func (s *lexSpec) KindIDAndName(mode ModeID, modeKind ModeKindID) (KindID, string) {
	id := s.kindIDs[mode][modeKind]
	return id, s.kindNames[id]
}
`

func genTemplateFuncs(clspec *spec.CompiledLexSpec) template.FuncMap {
	fns := template.FuncMap{
		"genPopTable": func() string {
			var b strings.Builder
			fmt.Fprintf(&b, "[][]bool{\n")
			for i, s := range clspec.Specs {
				if i == spec.LexModeIDNil.Int() {
					fmt.Fprintf(&b, "nil,\n")
					continue
				}

				fmt.Fprintf(&b, "{")
				fmt.Fprintf(&b, "%v", s.Pop[0] != 0)
				for _, v := range s.Pop[1:] {
					fmt.Fprintf(&b, ", %v", v != 0)
				}
				fmt.Fprintf(&b, "},\n")
			}
			fmt.Fprintf(&b, "}")
			return b.String()
		},
		"genPushTable": func() string {
			var b strings.Builder
			fmt.Fprintf(&b, "[][]ModeID{\n")
			for i, s := range clspec.Specs {
				if i == spec.LexModeIDNil.Int() {
					fmt.Fprintf(&b, "nil,\n")
					continue
				}

				fmt.Fprintf(&b, "{")
				fmt.Fprintf(&b, "%v", s.Push[0])
				for _, v := range s.Push[1:] {
					fmt.Fprintf(&b, ", %v", v)
				}
				fmt.Fprintf(&b, "},\n")
			}
			fmt.Fprintf(&b, "}")
			return b.String()
		},
		"genModeNameTable": func() string {
			var b strings.Builder
			fmt.Fprintf(&b, "[]string{\n")
			for i, name := range clspec.ModeNames {
				if i == spec.LexModeIDNil.Int() {
					fmt.Fprintf(&b, "%#v,\n", "")
					continue
				}

				fmt.Fprintf(&b, "%#v,\n", name)
			}
			fmt.Fprintf(&b, "}")
			return b.String()
		},
		"genInitialStateTable": func() string {
			var b strings.Builder
			fmt.Fprintf(&b, "[]StateID{\n")
			for i, s := range clspec.Specs {
				if i == spec.LexModeIDNil.Int() {
					fmt.Fprintf(&b, "%v,\n", spec.StateIDNil)
					continue
				}

				fmt.Fprintf(&b, "%v,\n", s.DFA.InitialStateID)
			}
			fmt.Fprintf(&b, "}")
			return b.String()
		},
		"genAcceptTable": func() string {
			var b strings.Builder
			fmt.Fprintf(&b, "[][]ModeKindID{\n")
			for i, s := range clspec.Specs {
				if i == spec.LexModeIDNil.Int() {
					fmt.Fprintf(&b, "nil,\n")
					continue
				}

				fmt.Fprintf(&b, "{")
				fmt.Fprintf(&b, "%v", s.DFA.AcceptingStates[0])
				for _, v := range s.DFA.AcceptingStates[1:] {
					fmt.Fprintf(&b, ", %v", v)
				}
				fmt.Fprintf(&b, "},\n")
			}
			fmt.Fprintf(&b, "}")
			return b.String()
		},
		"genKindIDTable": func() string {
			var b strings.Builder
			fmt.Fprintf(&b, "[][]KindID{\n")
			for i, ids := range clspec.KindIDs {
				if i == spec.LexModeIDNil.Int() {
					fmt.Fprintf(&b, "nil,\n")
					continue
				}

				fmt.Fprintf(&b, "{")
				fmt.Fprintf(&b, "%v", ids[0])
				for _, v := range ids[1:] {
					fmt.Fprintf(&b, ", %v", v)
				}
				fmt.Fprintf(&b, "},\n")
			}
			fmt.Fprintf(&b, "}")
			return b.String()
		},
		"genKindNameTable": func() string {
			var b strings.Builder
			fmt.Fprintf(&b, "[]string{\n")
			for i, name := range clspec.KindNames {
				if i == spec.LexKindIDNil.Int() {
					fmt.Fprintf(&b, "%#v,\n", "")
					continue
				}

				fmt.Fprintf(&b, "%#v,\n", name)
			}
			fmt.Fprintf(&b, "}")
			return b.String()
		},
	}

	switch clspec.CompressionLevel {
	case 2:
		fns["genRowNums"] = func() string {
			var b strings.Builder
			fmt.Fprintf(&b, "[][]int{\n")
			for i, s := range clspec.Specs {
				if i == spec.LexModeIDNil.Int() {
					fmt.Fprintf(&b, "nil,\n")
					continue
				}

				fmt.Fprintf(&b, "{")
				fmt.Fprintf(&b, "%v", s.DFA.Transition.RowNums[0])
				for _, v := range s.DFA.Transition.RowNums[1:] {
					fmt.Fprintf(&b, ", %v", v)
				}
				fmt.Fprintf(&b, "},\n")
			}
			fmt.Fprintf(&b, "}")
			return b.String()
		}

		fns["genRowDisplacements"] = func() string {
			var b strings.Builder
			fmt.Fprintf(&b, "[][]int{\n")
			for i, s := range clspec.Specs {
				if i == spec.LexModeIDNil.Int() {
					fmt.Fprintf(&b, "nil,\n")
					continue
				}

				fmt.Fprintf(&b, "{")
				fmt.Fprintf(&b, "%v", s.DFA.Transition.UniqueEntries.RowDisplacement[0])
				for _, d := range s.DFA.Transition.UniqueEntries.RowDisplacement[1:] {
					fmt.Fprintf(&b, ", %v", d)
				}
				fmt.Fprintf(&b, "},\n")
			}
			fmt.Fprintf(&b, "}")
			return b.String()
		}

		fns["genBounds"] = func() string {
			var b strings.Builder
			fmt.Fprintf(&b, "[][]int{\n")
			for i, s := range clspec.Specs {
				if i == spec.LexModeIDNil.Int() {
					fmt.Fprintf(&b, "nil,\n")
					continue
				}

				fmt.Fprintf(&b, "{")
				fmt.Fprintf(&b, "%v", s.DFA.Transition.UniqueEntries.Bounds[0])
				for _, v := range s.DFA.Transition.UniqueEntries.Bounds[1:] {
					fmt.Fprintf(&b, ", %v", v)
				}
				fmt.Fprintf(&b, "},\n")
			}
			fmt.Fprintf(&b, "}")
			return b.String()
		}

		fns["genEntries"] = func() string {
			var b strings.Builder
			fmt.Fprintf(&b, "[][]StateID{\n")
			for i, s := range clspec.Specs {
				if i == spec.LexModeIDNil.Int() {
					fmt.Fprintf(&b, "nil,\n")
					continue
				}

				fmt.Fprintf(&b, "{")
				fmt.Fprintf(&b, "%v", s.DFA.Transition.UniqueEntries.Entries[0])
				for _, v := range s.DFA.Transition.UniqueEntries.Entries[1:] {
					fmt.Fprintf(&b, ", %v", v)
				}
				fmt.Fprintf(&b, "},\n")
			}
			fmt.Fprintf(&b, "}")
			return b.String()
		}

		fns["genOriginalColCounts"] = func() string {
			return "nil"
		}
	case 1:
		fns["genRowNums"] = func() string {
			var b strings.Builder
			fmt.Fprintf(&b, "[][]int{\n")
			for i, s := range clspec.Specs {
				if i == spec.LexModeIDNil.Int() {
					fmt.Fprintf(&b, "nil,\n")
					continue
				}

				fmt.Fprintf(&b, "{")
				fmt.Fprintf(&b, "%v", s.DFA.Transition.RowNums[0])
				for _, v := range s.DFA.Transition.RowNums[1:] {
					fmt.Fprintf(&b, ", %v", v)
				}
				fmt.Fprintf(&b, "},\n")
			}
			fmt.Fprintf(&b, "}")
			return b.String()
		}

		fns["genRowDisplacements"] = func() string {
			return "nil"
		}

		fns["genBounds"] = func() string {
			return "nil"
		}

		fns["genEntries"] = func() string {
			var b strings.Builder
			fmt.Fprintf(&b, "[][]StateID{\n")
			for i, s := range clspec.Specs {
				if i == spec.LexModeIDNil.Int() {
					fmt.Fprintf(&b, "nil,\n")
					continue
				}

				fmt.Fprintf(&b, "{")
				fmt.Fprintf(&b, "%v", s.DFA.Transition.UncompressedUniqueEntries[0])
				for _, v := range s.DFA.Transition.UncompressedUniqueEntries[1:] {
					fmt.Fprintf(&b, ", %v", v)
				}
				fmt.Fprintf(&b, "},\n")
			}
			fmt.Fprintf(&b, "}")
			return b.String()
		}

		fns["genOriginalColCounts"] = func() string {
			var b strings.Builder
			fmt.Fprintf(&b, "[]int{\n")
			for i, s := range clspec.Specs {
				if i == spec.LexModeIDNil.Int() {
					fmt.Fprintf(&b, "0,\n")
					continue
				}

				fmt.Fprintf(&b, "%v,\n", s.DFA.Transition.OriginalColCount)
			}
			fmt.Fprintf(&b, "}")
			return b.String()
		}
	default:
		fns["genRowNums"] = func() string {
			return "nil"
		}

		fns["genRowDisplacements"] = func() string {
			return "nil"
		}

		fns["genBounds"] = func() string {
			return "nil"
		}

		fns["genEntries"] = func() string {
			var b strings.Builder
			fmt.Fprintf(&b, "[][]StateID{\n")
			for i, s := range clspec.Specs {
				if i == spec.LexModeIDNil.Int() {
					fmt.Fprintf(&b, "nil,\n")
					continue
				}

				fmt.Fprintf(&b, "{")
				fmt.Fprintf(&b, "%v", s.DFA.UncompressedTransition[0])
				for _, v := range s.DFA.UncompressedTransition[1:] {
					fmt.Fprintf(&b, ", %v", v)
				}
				fmt.Fprintf(&b, "},\n")
			}
			fmt.Fprintf(&b, "}")
			return b.String()
		}

		fns["genOriginalColCounts"] = func() string {
			var b strings.Builder
			fmt.Fprintf(&b, "[]int{\n")
			for i, s := range clspec.Specs {
				if i == spec.LexModeIDNil.Int() {
					fmt.Fprintf(&b, "0,\n")
					continue
				}

				fmt.Fprintf(&b, "%v,\n", s.DFA.ColCount)
			}
			fmt.Fprintf(&b, "}")
			return b.String()
		}
	}

	return fns
}