// gen generates contracts/counter/counter.wasm — binary WASM for the counter contract.
// Run from the repo root: go run ./contracts/counter/gen/
//
// Contract methods exported: increment, get, reset
// Host imports from "env": put_u64, get_u64, log, get_caller, get_state, set_state
package main

import (
	"fmt"
	"os"
)

// ── LEB128 ───────────────────────────────────────────────────────────────────

func u(v uint64) []byte {
	var b []byte
	for {
		bt := byte(v & 0x7f)
		v >>= 7
		if v != 0 {
			bt |= 0x80
		}
		b = append(b, bt)
		if v == 0 {
			return b
		}
	}
}

func s(v int64) []byte {
	var b []byte
	for {
		bt := byte(v & 0x7f)
		v >>= 7
		sign := (bt & 0x40) != 0
		if (v == 0 && !sign) || (v == -1 && sign) {
			return append(b, bt)
		}
		b = append(b, bt|0x80)
	}
}

// ── Builders ──────────────────────────────────────────────────────────────────

func cat(slices ...[]byte) []byte {
	var out []byte
	for _, s := range slices {
		out = append(out, s...)
	}
	return out
}

func wstr(str string) []byte { return cat(u(uint64(len(str))), []byte(str)) }

func section(id byte, content []byte) []byte {
	return cat([]byte{id}, u(uint64(len(content))), content)
}

// vec encodes a vector: count followed by concatenated items.
func vec(items ...[]byte) []byte {
	out := u(uint64(len(items)))
	for _, it := range items {
		out = append(out, it...)
	}
	return out
}

// functype encodes a WASM function type (0x60 prefix).
func functype(params, results []byte) []byte {
	return cat([]byte{0x60}, u(uint64(len(params))), params, u(uint64(len(results))), results)
}

// importFunc encodes a function import entry.
func importFunc(mod, name string, typeIdx uint32) []byte {
	return cat(wstr(mod), wstr(name), []byte{0x00}, u(uint64(typeIdx)))
}

// exportEntry encodes an export entry.
func exportEntry(name string, kind byte, idx uint32) []byte {
	return cat(wstr(name), []byte{kind}, u(uint64(idx)))
}

// dataSegment encodes an active data segment for memory 0.
func dataSegment(offset int32, data []byte) []byte {
	return cat(
		[]byte{0x00},             // active segment, implicit mem 0
		[]byte{0x41}, s(int64(offset)), []byte{0x0B}, // i32.const offset; end
		u(uint64(len(data))), data,
	)
}

// funcBody encodes one function body: localDecls + instructions + end.
func funcBody(localDecls []byte, instrs ...[]byte) []byte {
	inner := cat(localDecls)
	for _, ins := range instrs {
		inner = append(inner, ins...)
	}
	inner = append(inner, 0x0B) // end
	return cat(u(uint64(len(inner))), inner)
}

// noLocals is an empty local decl list.
var noLocals = u(0)

// localDecl encodes n locals of a given type.
func localDecl(n uint32, typ byte) []byte { return cat(u(uint64(n)), []byte{typ}) }
func withLocals(decls ...[]byte) []byte {
	return cat(u(uint64(len(decls))), cat(decls...))
}

// ── Instructions ──────────────────────────────────────────────────────────────

const (
	tI32 byte = 0x7F
	tI64 byte = 0x7E
)

func call(fn uint32) []byte       { return cat([]byte{0x10}, u(uint64(fn))) }
func lget(i uint32) []byte        { return cat([]byte{0x20}, u(uint64(i))) }
func lset(i uint32) []byte        { return cat([]byte{0x21}, u(uint64(i))) }
func ic32(v int32) []byte         { return cat([]byte{0x41}, s(int64(v))) }
func ic64(v int64) []byte         { return cat([]byte{0x42}, s(v)) }
func block_() []byte              { return []byte{0x02, 0x40} }
func loop_() []byte               { return []byte{0x03, 0x40} }
func if_() []byte                 { return []byte{0x04, 0x40} }
func end_() []byte                { return []byte{0x0B} }
func br_(lbl uint32) []byte       { return cat([]byte{0x0C}, u(uint64(lbl))) }
func brIf_(lbl uint32) []byte     { return cat([]byte{0x0D}, u(uint64(lbl))) }
func return_() []byte             { return []byte{0x0F} }
func i32Eqz() []byte              { return []byte{0x45} }
func i32Ne() []byte               { return []byte{0x47} }
func i32GeU() []byte              { return []byte{0x4F} }
func i32Add() []byte              { return []byte{0x6A} }
func i64Add() []byte              { return []byte{0x7C} }
func i32Load8U() []byte           { return []byte{0x2D, 0x00, 0x00} } // align=0, offset=0

// ── Memory layout constants ───────────────────────────────────────────────────

const (
	offCounter     = 0x00 // "counter" (7 bytes)
	offOwner       = 0x10 // "owner"   (5 bytes)
	offIncMsg      = 0x20 // "incremented" (11 bytes)
	offGetMsg      = 0x30 // "get called" (10 bytes)
	offResetOk     = 0x40 // "reset ok" (8 bytes)
	offUnauth      = 0x50 // "unauthorized" (12 bytes)
	offCallerBuf   = 0x60 // caller buf (128 bytes)
	offOwnerBuf    = 0xE0 // owner buf  (128 bytes)
)

// Import function indices
const (
	fnPutU64    = 0 // put_u64(keyPtr, keyLen i32, val i64)
	fnGetU64    = 1 // get_u64(keyPtr, keyLen i32) → i64
	fnLog       = 2 // log(msgPtr, msgLen i32)
	fnGetCaller = 3 // get_caller(bufPtr, bufLen i32) → i32
	fnGetState  = 4 // get_state(kPtr,kLen,dPtr,dLen i32) → i32
	fnSetState  = 5 // set_state(kPtr,kLen,vPtr,vLen i32)
)

// Local function indices (imports are 0-5, locals start at 6)
const (
	fnIncrement = 6
	fnGet       = 7
	fnReset     = 8
)

func main() {
	// ── Type section ─────────────────────────────────────────────────────────
	// Type 0: (i32,i32,i64)→()      put_u64
	// Type 1: (i32,i32)→(i64)       get_u64
	// Type 2: (i32,i32)→()          log
	// Type 3: (i32,i32)→(i32)       get_caller
	// Type 4: (i32,i32,i32,i32)→(i32) get_state
	// Type 5: (i32,i32,i32,i32)→()  set_state
	// Type 6: ()→()                  increment, get, reset
	typeSection := section(0x01, vec(
		functype([]byte{tI32, tI32, tI64}, []byte{}),               // 0
		functype([]byte{tI32, tI32}, []byte{tI64}),                 // 1
		functype([]byte{tI32, tI32}, []byte{}),                     // 2
		functype([]byte{tI32, tI32}, []byte{tI32}),                 // 3
		functype([]byte{tI32, tI32, tI32, tI32}, []byte{tI32}),    // 4
		functype([]byte{tI32, tI32, tI32, tI32}, []byte{}),        // 5
		functype([]byte{}, []byte{}),                               // 6
	))

	// ── Import section ────────────────────────────────────────────────────────
	importSection := section(0x02, vec(
		importFunc("env", "put_u64",    fnPutU64),
		importFunc("env", "get_u64",    fnGetU64),
		importFunc("env", "log",        fnLog),
		importFunc("env", "get_caller", fnGetCaller),
		importFunc("env", "get_state",  fnGetState),
		importFunc("env", "set_state",  fnSetState),
	))

	// ── Function section: 3 local functions, all type 6 ──────────────────────
	functionSection := section(0x03, vec(u(6), u(6), u(6)))

	// ── Memory section: 1 page (64 KiB) ──────────────────────────────────────
	// limits type 0x00 = min only; type 0x01 = min+max
	memorySection := section(0x05, vec(cat([]byte{0x00}, u(1)))) // min=1, no max

	// ── Export section ────────────────────────────────────────────────────────
	exportSection := section(0x07, vec(
		exportEntry("memory",    0x02, 0),
		exportEntry("increment", 0x00, fnIncrement),
		exportEntry("get",       0x00, fnGet),
		exportEntry("reset",     0x00, fnReset),
	))

	// ── Data section ──────────────────────────────────────────────────────────
	dataSection := section(0x0B, cat(
		u(6), // 6 segments
		dataSegment(offCounter,   []byte("counter")),
		dataSegment(offOwner,     []byte("owner")),
		dataSegment(offIncMsg,    []byte("incremented")),
		dataSegment(offGetMsg,    []byte("get called")),
		dataSegment(offResetOk,   []byte("reset ok")),
		dataSegment(offUnauth,    []byte("unauthorized")),
	))

	// ── Code section ─────────────────────────────────────────────────────────

	// increment():
	//   local $val i64
	//   $val = get_u64("counter")
	//   $val++
	//   put_u64("counter", $val)
	//   log("incremented")
	incrementBody := funcBody(
		withLocals(localDecl(1, tI64)),
		ic32(offCounter), ic32(7), call(fnGetU64), lset(0),
		lget(0), ic64(1), i64Add(), lset(0),
		ic32(offCounter), ic32(7), lget(0), call(fnPutU64),
		ic32(offIncMsg), ic32(11), call(fnLog),
	)

	// get():
	//   log("get called")
	getBody := funcBody(
		noLocals,
		ic32(offGetMsg), ic32(10), call(fnLog),
	)

	// reset():
	//   locals: callerLen(0), ownerLen(1), i(2), same(3)  — all i32
	//   callerLen = get_caller(callerBuf, 128)
	//   ownerLen  = get_state("owner", ownerBuf, 128)
	//   if ownerLen == 0:
	//     set_state("owner", callerBuf[:callerLen])
	//     put_u64("counter", 0)
	//     log("reset ok")
	//     return
	//   if callerLen != ownerLen: log unauthorized; return
	//   same = 1; i = 0
	//   block:
	//     loop:
	//       if i >= callerLen: br 1 (exit block)
	//       if callerBuf[i] != ownerBuf[i]: same=0; br 1
	//       i++; continue loop
	//   if !same: log unauthorized; return
	//   put_u64("counter", 0); log("reset ok")
	resetBody := funcBody(
		withLocals(localDecl(4, tI32)),
		// callerLen = get_caller(callerBuf, 128)
		ic32(offCallerBuf), ic32(128), call(fnGetCaller), lset(0),
		// ownerLen = get_state("owner", 5, ownerBuf, 128)
		ic32(offOwner), ic32(5), ic32(offOwnerBuf), ic32(128), call(fnGetState), lset(1),
		// if ownerLen == 0:
		lget(1), i32Eqz(), if_(),
		  ic32(offOwner), ic32(5), ic32(offCallerBuf), lget(0), call(fnSetState),
		  ic32(offCounter), ic32(7), ic64(0), call(fnPutU64),
		  ic32(offResetOk), ic32(8), call(fnLog),
		  return_(),
		end_(),
		// if callerLen != ownerLen: unauthorized
		lget(0), lget(1), i32Ne(), if_(),
		  ic32(offUnauth), ic32(12), call(fnLog),
		  return_(),
		end_(),
		// same = 1; i = 0
		ic32(1), lset(3),
		ic32(0), lset(2),
		// block $break
		block_(),
		  loop_(),
		    lget(2), lget(0), i32GeU(), brIf_(1), // i >= callerLen → break
		    // load callerBuf[i]
		    ic32(offCallerBuf), lget(2), i32Add(), i32Load8U(),
		    // load ownerBuf[i]
		    ic32(offOwnerBuf),  lget(2), i32Add(), i32Load8U(),
		    i32Ne(), if_(),
		      ic32(0), lset(3),
		      br_(2), // break out of block
		    end_(),
		    lget(2), ic32(1), i32Add(), lset(2),
		    br_(0), // continue loop
		  end_(),
		end_(),
		// if !same: unauthorized
		lget(3), i32Eqz(), if_(),
		  ic32(offUnauth), ic32(12), call(fnLog),
		  return_(),
		end_(),
		// authorized
		ic32(offCounter), ic32(7), ic64(0), call(fnPutU64),
		ic32(offResetOk), ic32(8), call(fnLog),
	)

	codeSection := section(0x0A, cat(u(3), incrementBody, getBody, resetBody))

	// ── Assemble module ───────────────────────────────────────────────────────
	module := cat(
		[]byte{0x00, 0x61, 0x73, 0x6d}, // magic \0asm
		[]byte{0x01, 0x00, 0x00, 0x00}, // version 1
		typeSection,
		importSection,
		functionSection,
		memorySection,
		exportSection,
		dataSection,
		codeSection,
	)

	out := "contracts/counter/counter.wasm"
	if err := os.WriteFile(out, module, 0644); err != nil {
		fmt.Fprintln(os.Stderr, "write:", err)
		os.Exit(1)
	}
	fmt.Printf("Written %s (%d bytes)\n", out, len(module))
}