package preproc

import (
	"strings"
	"unicode/utf8"

	"github.com/armon/go-radix"
)

type DefineList struct {
	tree *radix.Tree
}

func NewDefineList() *DefineList { return &DefineList{tree: radix.New()} }

// Add inserts or updates a define->value mapping.
func (d *DefineList) Add(define, value string) {
	d.tree.Insert(define, expandDollarEscapes(value))
}

// Delete removes a define. Returns true if it existed.
func (d *DefineList) Delete(define string) bool {
	_, ok := d.tree.Delete(define)
	return ok
}

// Defined reports whether a define exists.
func (d *DefineList) Defined(s string) bool {
	_, ok := d.tree.Get(s)
	return ok
}

// typed wrapper to avoid type assertions sprinkled in hot code
func (d *DefineList) longestPrefixString(s string) (key string, val string, ok bool) {
	k, v, ok := d.tree.LongestPrefix(s)
	if !ok {
		return "", "", false
	}
	return k, v.(string), true
}

// ReplaceDefines performs a single-pass, longest-prefix replacement.
// Case-sensitive. Matches anywhere. No recursive re-expansion.
func (d *DefineList) ReplaceDefines(s string) string {
	var b strings.Builder
	b.Grow(len(s)) // baseline; grows if expansions occur

	for i := 0; i < len(s); {
		if key, val, ok := d.longestPrefixString(s[i:]); ok && key != "" {
			b.WriteString(val)
			i += len(key) // advance by matched bytes (UTF-8-safe)
			continue
		}
		_, size := utf8.DecodeRuneInString(s[i:])
		b.WriteString(s[i : i+size])
		i += size
	}
	return b.String()
}

func isHex(c byte) bool {
	return (c >= '0' && c <= '9') || (c >= 'A' && c <= 'F') || (c >= 'a' && c <= 'f')
}

func hexToByte(hi, lo byte) byte {
	return hexDigit(hi)<<4 | hexDigit(lo)
}

func hexDigit(c byte) byte {
	if c >= '0' && c <= '9' {
		return c - '0'
	}
	if c >= 'A' && c <= 'F' {
		return c - 'A' + 10
	}
	if c >= 'a' && c <= 'f' {
		return c - 'a' + 10
	}
	return 0
}

func expandDollarEscapes(s string) string {
	var b strings.Builder
	for i := 0; i < len(s); i++ {
		if s[i] == '$' {
			if i+1 < len(s) && s[i+1] == '$' {
				// $$ -> $
				b.WriteByte('$')
				i++ // skip second $
			} else if i+2 < len(s) && isHex(s[i+1]) && isHex(s[i+2]) {
				// $XX -> character
				val := hexToByte(s[i+1], s[i+2])
				b.WriteByte(val)
				i += 2 // skip hex digits
			} else {
				b.WriteByte('$')
			}
		} else {
			b.WriteByte(s[i])
		}
	}
	return b.String()
}