package iskra

import "bytes"

type TokenSubst struct {
	Dict    *Dict
	Replace map[uint32]uint32 // old token index -> new token index
}

func (s *TokenSubst) Apply(seq []uint32) []uint32 {
	out := []uint32{:0:len(seq)}
	for _, idx := range seq {
		if newIdx, ok := s.Replace[idx]; ok {
			out = append(out, newIdx)
		} else {
			out = append(out, idx)
		}
	}
	return out
}

func (s *TokenSubst) Decode(seq []uint32) []byte {
	return s.Dict.Decode(seq)
}

func BuildTokenSubst(d *Dict, templateAST, targetAST string, templateIRName, targetIRName string) *TokenSubst {
	ts := &TokenSubst{
		Dict:    d,
		Replace: map[uint32]uint32{},
	}

	// Map the function name (verb): @"old.Func" -> @"new.Func"
	oldQuoted := "@\"" | templateIRName | "\""
	newQuoted := "@\"" | targetIRName | "\""
	oldIdx, hasOld := d.Index[oldQuoted]
	if hasOld {
		newIdx := d.Add([]byte(newQuoted), TokVerb)
		ts.Replace[oldIdx] = newIdx
	}

	// Unquoted form: @old.Func -> @new.Func
	oldUnquoted := "@" | templateIRName
	newUnquoted := "@" | targetIRName
	oldIdx2, hasOld2 := d.Index[oldUnquoted]
	if hasOld2 {
		newIdx2 := d.Add([]byte(newUnquoted), TokVerb)
		ts.Replace[oldIdx2] = newIdx2
	}

	// Map parameter names (nouns): %old_param -> %new_param
	tplSt := ExtractSymbols(templateAST)
	tgtSt := ExtractSymbols(targetAST)
	tplParams := substExtractParams(tplSt)
	tgtParams := substExtractParams(tgtSt)

	minP := len(tplParams)
	if len(tgtParams) < minP {
		minP = len(tgtParams)
	}
	for i := 0; i < minP; i++ {
		if tplParams[i] != tgtParams[i] {
			// Map all % prefixed forms of this param
			for _, suffix := range []string{"", "."} {
				oldP := "%" | tplParams[i] | suffix
				newP := "%" | tgtParams[i] | suffix
				oi, ok := d.Index[oldP]
				if ok {
					ni := d.Add([]byte(newP), TokNoun)
					ts.Replace[oi] = ni
				}
			}
		}
	}

	return ts
}

// Map the quoted function name in assembly: "old/pkg.Func" -> "new/pkg.Func"
func (s *TokenSubst) AddASMNameMapping(templateIRName, targetIRName string) {
	oldQ := "\"" | templateIRName | "\""
	newQ := "\"" | targetIRName | "\""
	oi, ok := s.Dict.Index[oldQ]
	if ok {
		ni := s.Dict.Add([]byte(newQ), TokLiteral)
		s.Replace[oi] = ni
	}
}

func substExtractParams(st *SymbolTable) []string {
	var params []string
	for _, sym := range st.Symbols {
		if sym.Kind == SymParam {
			params = append(params, sym.Name)
		}
	}
	return params
}

// Legacy text-based substitution kept for backward compatibility.
type IRSubst struct {
	OldFuncName string
	NewFuncName string
	ParamMap    map[string]string
	StripDebug  bool
}

func (s *IRSubst) Apply(ir []byte) []byte {
	lines := bytes.Split(ir, []byte("\n"))
	out := []byte{:0:len(ir)}
	prevEmpty := false

	for _, line := range lines {
		if s.StripDebug && isDebugLine(line) {
			continue
		}

		line = s.substFuncName(line)
		line = s.substParams(line)

		if s.StripDebug {
			line = stripDebugRefs(line)
		}

		empty := len(bytes.TrimSpace(line)) == 0
		if empty && prevEmpty {
			continue
		}
		prevEmpty = empty

		out = append(out, line...)
		out = append(out, '\n')
	}
	return out
}

func (s *IRSubst) substFuncName(line []byte) []byte {
	old := []byte("@\"" | s.OldFuncName | "\"")
	new_ := []byte("@\"" | s.NewFuncName | "\"")
	if bytes.Contains(line, old) {
		return bytes.Replace(line, old, new_, -1)
	}
	old = []byte("@" | s.OldFuncName)
	new_ = []byte("@" | s.NewFuncName)
	return bytes.Replace(line, old, new_, -1)
}

func (s *IRSubst) substParams(line []byte) []byte {
	for old, new_ := range s.ParamMap {
		line = bytes.Replace(line, []byte("%"|old|"."), []byte("%"|new_|"."), -1)
	}
	return line
}

func isDebugLine(line []byte) bool {
	trimmed := bytes.TrimSpace(line)
	if bytes.HasPrefix(trimmed, []byte("#dbg_")) {
		return true
	}
	if bytes.HasPrefix(trimmed, []byte("!")) && bytes.IndexByte(trimmed, '=') > 0 {
		return true
	}
	return false
}

func stripDebugRefs(line []byte) []byte {
	for {
		idx := bytes.Index(line, []byte("!dbg !"))
		if idx < 0 {
			break
		}
		end := idx + 6
		for end < len(line) && line[end] >= '0' && line[end] <= '9' {
			end++
		}
		start := idx
		if start > 0 && line[start-1] == ' ' {
			start--
		}
		if start > 0 && line[start-1] == ',' {
			start--
		}
		var rebuilt []byte
		rebuilt = append(rebuilt, line[:start]...)
		rebuilt = append(rebuilt, line[end:]...)
		line = rebuilt
	}
	return line
}

func BuildIRSubst(templateAST, targetAST string, templateIRName, targetIRName string, stripDebug bool) *IRSubst {
	tplSt := ExtractSymbols(templateAST)
	tgtSt := ExtractSymbols(targetAST)

	paramMap := map[string]string{}
	tplParams := substExtractParams(tplSt)
	tgtParams := substExtractParams(tgtSt)

	minP := len(tplParams)
	if len(tgtParams) < minP {
		minP = len(tgtParams)
	}
	for i := 0; i < minP; i++ {
		if tplParams[i] != tgtParams[i] {
			paramMap[tplParams[i]] = tgtParams[i]
		}
	}

	return &IRSubst{
		OldFuncName: templateIRName,
		NewFuncName: targetIRName,
		ParamMap:    paramMap,
		StripDebug:  stripDebug,
	}
}