// Bootstrap Go enzyme organ: tokenizes Go source code into typed elements.
//
// This is the hand-written template for the first enzyme organ. The organism
// will eventually learn to modify and improve it, but this provides the
// initial capability to digest Go source.
//
// Opcodes:
//   0x01 CanDigest(sample) -> 0x01 (yes) or 0x00 (no)
//   0x02 Digest(source)    -> JSON array of {type, value} pairs
//
// Compile with:
//   GOOS=wasip1 GOARCH=wasm go build -buildmode=c-shared -o enzyme_go.wasm .
package main

import (
	"encoding/json"
	"go/scanner"
	"go/token"
	"unsafe"
)

// Memory pool: bump allocator for guest-side allocations.
var pool [1 << 20]byte // 1MB
var poolOffset uint32

//go:wasmexport alloc
func alloc(size uint32) uint32 {
	if poolOffset+size > uint32(len(pool)) {
		poolOffset = 0
	}
	ptr := poolOffset
	poolOffset += size
	return uint32(uintptr(unsafe.Pointer(&pool[ptr])))
}

//go:wasmexport dealloc
func dealloc(ptr uint32, size uint32) {}

// Element is a typed lattice element produced by digestion.
type Element struct {
	Type  string `json:"type"`
	Value string `json:"value"`
}

//go:wasmexport process
func process(inputPtr uint32, inputLen uint32) uint64 {
	base := uint32(uintptr(unsafe.Pointer(&pool[0])))
	offset := inputPtr - base

	if offset+inputLen > uint32(len(pool)) {
		return 0
	}

	input := pool[offset : offset+inputLen]
	if len(input) == 0 {
		return 0
	}

	opcode := input[0]
	data := input[1:]

	var result []byte

	switch opcode {
	case 0x01: // CanDigest: check if input looks like Go source.
		result = []byte{0x00}
		for i := 0; i+8 <= len(data); i++ {
			if string(data[i:i+8]) == "package " {
				result = []byte{0x01}
				break
			}
		}

	case 0x02: // Digest: tokenize Go source into elements.
		elements := tokenize(data)
		result, _ = json.Marshal(elements)
		if result == nil {
			result = []byte("[]")
		}

	default:
		result = []byte("unknown opcode")
	}

	outPtr := alloc(uint32(len(result)))
	outOffset := outPtr - base
	copy(pool[outOffset:], result)

	return uint64(outPtr)<<32 | uint64(len(result))
}

// tokenize uses go/scanner to break Go source into typed elements.
func tokenize(src []byte) []Element {
	fset := token.NewFileSet()
	file := fset.AddFile("input.go", fset.Base(), len(src))

	var s scanner.Scanner
	s.Init(file, src, nil, scanner.ScanComments)

	var elements []Element

	for {
		_, tok, lit := s.Scan()
		if tok == token.EOF {
			break
		}

		elem := Element{
			Type:  classifyToken(tok),
			Value: lit,
		}
		if lit == "" {
			elem.Value = tok.String()
		}

		elements = append(elements, elem)
	}

	return elements
}

// classifyToken maps Go token types to lattice element types.
func classifyToken(tok token.Token) string {
	switch {
	case tok == token.IDENT:
		return "ident"
	case tok == token.INT || tok == token.FLOAT || tok == token.IMAG ||
		tok == token.CHAR || tok == token.STRING:
		return "literal"
	case tok == token.COMMENT:
		return "comment"
	case tok == token.PACKAGE:
		return "package"
	case tok == token.IMPORT:
		return "import"
	case tok == token.FUNC:
		return "func"
	case tok == token.TYPE:
		return "type"
	case tok == token.STRUCT:
		return "struct"
	case tok == token.INTERFACE:
		return "interface"
	case tok == token.RETURN:
		return "return"
	case tok == token.IF:
		return "if"
	case tok == token.FOR:
		return "for"
	case tok == token.SWITCH:
		return "switch"
	case tok == token.SELECT:
		return "select"
	case tok == token.GO:
		return "go"
	case tok == token.CHAN:
		return "chan"
	case tok == token.VAR || tok == token.CONST:
		return "var"
	case tok.IsOperator():
		return "op"
	default:
		return "punct"
	}
}

func main() {}