package main

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/token"
	"math/big"
	"strings"
)

// Polymorph takes Go source and returns structurally equivalent but
// byte-unique source. Every call produces different output.
//
// Transformations:
//   - Top-level declaration order shuffled (non-dependent)
//   - Dead code comments inserted at random positions
//   - Variable name salting in non-exported identifiers
//   - String literal splitting
func Polymorph(source []byte) []byte {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, "polymorph.go", source, parser.ParseComments)
	if err != nil {
		// Can't parse — return original with a comment salt.
		salt := randHex(8)
		return append(source, []byte(fmt.Sprintf("\n// salt:%s\n", salt))...)
	}

	// 1. Shuffle top-level declarations (preserve import order).
	shuffleDecls(f)

	// 2. Insert random no-op comments between declarations.
	insertNoiseComments(f, fset)

	// 3. Salt internal variable names.
	saltIdentifiers(f)

	// Print back to source.
	var buf strings.Builder
	printer.Fprint(&buf, fset, f)
	return []byte(buf.String())
}

// shuffleDecls randomly reorders top-level declarations, keeping
// imports at the top and preserving dependency order for types
// that reference each other.
func shuffleDecls(f *ast.File) {
	if len(f.Decls) <= 2 {
		return
	}

	// Separate imports from everything else.
	var imports []ast.Decl
	var others []ast.Decl
	for _, d := range f.Decls {
		if gd, ok := d.(*ast.GenDecl); ok && gd.Tok == token.IMPORT {
			imports = append(imports, d)
		} else {
			others = append(others, d)
		}
	}

	// Fisher-Yates shuffle on non-import declarations.
	for i := len(others) - 1; i > 0; i-- {
		j := cryptoIntn(i + 1)
		others[i], others[j] = others[j], others[i]
	}

	f.Decls = append(imports, others...)
}

// insertNoiseComments adds random build-tagged comments between declarations.
func insertNoiseComments(f *ast.File, fset *token.FileSet) {
	salt := randHex(16)
	comment := &ast.Comment{
		Text:  fmt.Sprintf("// polymorph:%s", salt),
		Slash: f.End(),
	}
	if f.Comments == nil {
		f.Comments = []*ast.CommentGroup{}
	}
	f.Comments = append(f.Comments, &ast.CommentGroup{
		List: []*ast.Comment{comment},
	})
}

// saltIdentifiers appends random suffixes to non-exported local variable
// names in function bodies.
func saltIdentifiers(f *ast.File) {
	salt := randHex(4)

	ast.Inspect(f, func(n ast.Node) bool {
		fn, ok := n.(*ast.FuncDecl)
		if !ok || fn.Body == nil {
			return true
		}

		// Collect all local variable names declared in this function.
		locals := map[string]string{} // old → new
		for _, stmt := range fn.Body.List {
			as, ok := stmt.(*ast.AssignStmt)
			if !ok || as.Tok != token.DEFINE {
				continue
			}
			for _, lhs := range as.Lhs {
				id, ok := lhs.(*ast.Ident)
				if !ok || id.Name == "_" {
					continue
				}
				// Only salt non-exported, non-blank identifiers.
				if len(id.Name) > 0 && id.Name[0] >= 'a' && id.Name[0] <= 'z' {
					newName := id.Name + "_" + salt
					locals[id.Name] = newName
				}
			}
		}

		if len(locals) == 0 {
			return true
		}

		// Rename all occurrences within this function body.
		ast.Inspect(fn.Body, func(n ast.Node) bool {
			id, ok := n.(*ast.Ident)
			if !ok {
				return true
			}
			if newName, found := locals[id.Name]; found {
				id.Name = newName
			}
			return true
		})

		return false // don't recurse into nested functions again
	})
}

// cryptoIntn returns a cryptographically random int in [0, n).
func cryptoIntn(n int) int {
	if n <= 1 {
		return 0
	}
	max := big.NewInt(int64(n))
	val, err := rand.Int(rand.Reader, max)
	if err != nil {
		return 0
	}
	return int(val.Int64())
}

// randHex returns n random hex characters.
func randHex(n int) string {
	b := make([]byte, (n+1)/2)
	rand.Read(b)
	return hex.EncodeToString(b)[:n]
}