// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: BSD-3-Clause

package asm

import (
	"fmt"
	"strings"

	"sourcedock.dev/petrbalvin/gasm-devkit/ast"
)

// Assemble encodes the body of a TEXT function into x86-64 machine code,
// resolving local labels to relative jump offsets and translating the FP/SP
// pseudo-registers onto the hardware stack pointer (matching the Go
// assembler's default frame-pointer behaviour).  Jumps start in the short
// (rel8) form and expand to rel32 when the settled displacement does not fit;
// sizes only grow, so the layout reaches a fixed point in a few passes.  CALL
// has no short form and is always rel32.
//
// Supported operands: registers, memory (real base register), immediates,
// FP/SP frame-relative operands, and local-label jumps.  SB (global symbol)
// operands require relocations and are not yet supported; the SIMD (VEX/AVX2)
// integer and shuffle/extract/permute/move set is in.
func Assemble(t *ast.Text) ([]byte, map[string]int, error) {
	code, _, labels, err := assemble(t, nil)
	return code, labels, err
}

// linkInfo carries file-level symbol context into a single-function assembly:
// the set of static symbols a GLOBL in the same file defines.  A nil link
// rejects SB operands outright (single-function assembly cannot resolve
// them).
type linkInfo struct {
	symbols map[string]bool
}

// sbPatch is a function-relative static-symbol relocation: the disp32 field
// at off must become the symbol's address minus after, where after is the
// function-relative address just past the instruction.
type sbPatch struct {
	off    int
	after  int
	name   string
	addend int64
}

// assemble encodes a TEXT body, returning the machine code, the static-symbol
// patch sites (for the file-level layout to resolve) and the label table.
func assemble(t *ast.Text, link *linkInfo) ([]byte, []sbPatch, map[string]int, error) {
	fi := computeFrame(t)
	chain := jumpChain(t)
	resolve := func(name string) string {
		if r, ok := chain[name]; ok {
			return r
		}
		return name
	}

	// Layout: iterate jump sizes to a fixed point.
	long := make([]bool, len(t.Body))
	sizes := make([]int, len(t.Body))
	offsets := map[string]int{}
	pcs := make([]int, len(t.Body))
	for {
		pos := len(fi.prologue)
		for i, stmt := range t.Body {
			switch s := stmt.(type) {
			case *ast.Label:
				offsets[s.Name.Text] = pos
			case *ast.Instr:
				sz, err := instrSize(s, fi, long[i], link)
				if err != nil {
					return nil, nil, nil, fmt.Errorf("%s: %w", s.Mnemonic.Text, err)
				}
				sizes[i] = sz
				pcs[i] = pos
				pos += sz
			}
		}
		// Expand any short jump whose displacement no longer fits rel8.
		changed := false
		for i, stmt := range t.Body {
			s, ok := stmt.(*ast.Instr)
			if !ok {
				continue
			}
			mnem := strings.ToUpper(s.Mnemonic.Text)
			if !isJumpMnemonic(mnem) || mnem == "CALL" || long[i] {
				continue
			}
			name, ok := labelName(s.Operands[0])
			if !ok {
				continue // reported during emission
			}
			target, ok := offsets[resolve(name)]
			if !ok {
				continue // reported during emission
			}
			rel := int64(target - (pcs[i] + jumpSize(mnem, false)))
			if !fits8(rel) {
				long[i] = true
				changed = true
			}
		}
		if !changed {
			break
		}
	}

	// Pass 2: emit.
	out := append([]byte(nil), fi.prologue...)
	var patches []sbPatch
	pos := len(fi.prologue)
	for i, stmt := range t.Body {
		s, ok := stmt.(*ast.Instr)
		if !ok {
			continue
		}
		code, ps, err := encodeInstr(s, pos, offsets, fi, long[i], resolve, link)
		if err != nil {
			return nil, nil, nil, fmt.Errorf("%s: %w", s.Mnemonic.Text, err)
		}
		if len(code) != sizes[i] {
			return nil, nil, nil, fmt.Errorf("%s: size mismatch (%d vs %d)", s.Mnemonic.Text, len(code), sizes[i])
		}
		patches = append(patches, ps...)
		out = append(out, code...)
		pos += len(code)
	}
	return out, patches, offsets, nil
}

// jumpChain precomputes jump-to-jump folding: a label whose first instruction
// is an unconditional local jump redirects its own jumpers to the ultimate
// target.  The Go toolchain chases exactly these chains (the linker's xfol
// pass) before it encodes branches, so matching its bytes requires the same
// redirection.
func jumpChain(t *ast.Text) map[string]string {
	// label → the target of its leading unconditional local JMP, if any.
	leadsTo := map[string]string{}
	for i, stmt := range t.Body {
		l, ok := stmt.(*ast.Label)
		if !ok {
			continue
		}
		// Stacked labels share an address: skip to the first instruction.
		j := i + 1
		for j < len(t.Body) {
			if _, isLabel := t.Body[j].(*ast.Label); !isLabel {
				break
			}
			j++
		}
		if j >= len(t.Body) {
			continue
		}
		in, ok := t.Body[j].(*ast.Instr)
		if !ok || strings.ToUpper(in.Mnemonic.Text) != "JMP" || len(in.Operands) != 1 {
			continue
		}
		if name, ok := labelName(in.Operands[0]); ok {
			leadsTo[l.Name.Text] = name
		}
	}
	// Chase each chain to its end, guarding against cycles.
	chain := map[string]string{}
	for name := range leadsTo {
		visited := map[string]bool{name: true}
		cur := name
		for {
			next, ok := leadsTo[cur]
			if !ok || visited[next] {
				break
			}
			visited[next] = true
			cur = next
		}
		if cur != name {
			chain[name] = cur
		}
	}
	return chain
}

// frameInfo carries the frame layout derived from the TEXT directive.
type frameInfo struct {
	size     int   // local frame size ($framesize)
	useFP    bool  // a frame pointer (BP) is set up
	fpAdjust int64 // added to x+N(FP) to reach the hardware SP-relative offset
	spAdjust int64 // x-N(SP) becomes (spAdjust - N)(SP)
	prologue []byte
	epilogue []byte
}

// computeFrame derives the frame layout, matching the Go assembler's default
// (a frame pointer is used whenever the function has a non-zero frame).
func computeFrame(t *ast.Text) frameInfo {
	fi := frameInfo{}
	if t.Frame != nil && t.Frame.Imm.HasVal {
		fi.size = int(t.Frame.Imm.Val)
	}
	if fi.size > 0 {
		fi.useFP = true
		fi.fpAdjust = int64(fi.size) + 16 // frame + saved BP + return address
		fi.spAdjust = int64(fi.size)
		fi.prologue = prologueBytes(fi.size)
		fi.epilogue = epilogueBytes(fi.size)
	} else {
		fi.fpAdjust = 8 // return address only
	}
	return fi
}

// prologueBytes emits: PUSHQ BP; MOVQ SP, BP; SUBQ $size, SP.
func prologueBytes(size int) []byte {
	out := []byte{0x55, 0x48, 0x89, 0xE5} // PUSHQ BP; MOVQ SP, BP
	return append(out, subSP(size)...)
}

// epilogueBytes emits: ADDQ $size, SP; POPQ BP.
func epilogueBytes(size int) []byte {
	out := addSP(size)
	return append(out, 0x5D) // POPQ BP
}

func subSP(size int) []byte { // SUBQ $size, SP
	if size >= -128 && size <= 127 {
		return []byte{0x48, 0x83, 0xEC, byte(int8(size))}
	}
	return append([]byte{0x48, 0x81, 0xEC}, le32(int64(size))...)
}

func addSP(size int) []byte { // ADDQ $size, SP
	if size >= -128 && size <= 127 {
		return []byte{0x48, 0x83, 0xC4, byte(int8(size))}
	}
	return append([]byte{0x48, 0x81, 0xC4}, le32(int64(size))...)
}

// instrSize returns the encoded length of an instruction (layout pass).
// encodeInstr already includes the epilogue for a RET in a frame-pointer
// function; jumps use their short or long form (never an epilogue).
func instrSize(s *ast.Instr, fi frameInfo, long bool, link *linkInfo) (int, error) {
	mnem := strings.ToUpper(s.Mnemonic.Text)
	if isJumpMnemonic(mnem) {
		return jumpSize(mnem, long), nil
	}
	code, _, err := encodeInstr(s, 0, nil, fi, false, nil, link)
	if err != nil {
		return 0, err
	}
	return len(code), nil
}

func isJumpMnemonic(mnem string) bool {
	if mnem == "JMP" || mnem == "CALL" {
		return true
	}
	_, ok := condCode(mnem)
	return ok
}

// jumpSize returns the length of a jump instruction in the requested form:
// short (rel8) where available, otherwise the rel32 form.  CALL is always
// rel32.
func jumpSize(mnem string, long bool) int {
	if mnem == "CALL" {
		return 5 // opcode + rel32
	}
	if !long {
		return 2 // opcode + rel8
	}
	if mnem == "JMP" {
		return 5 // E9 + rel32
	}
	return 6 // 0x0F 0x8x + rel32
}

// encodeInstr encodes one instruction, resolving jump targets against offsets
// (relative to pc, the instruction's own offset).  A RET in a frame-pointer
// function is prefixed with the epilogue.  resolve, when non-nil, redirects a
// jump label through the jump-to-jump chain before the offset lookup.
func encodeInstr(s *ast.Instr, pc int, offsets map[string]int, fi frameInfo, long bool, resolve func(string) string, link *linkInfo) ([]byte, []sbPatch, error) {
	mnem := strings.ToUpper(s.Mnemonic.Text)

	var prefix []byte
	if mnem == "RET" && fi.useFP {
		prefix = fi.epilogue
	}

	var code []byte
	var ps []sbPatch
	var err error
	if isJumpMnemonic(mnem) {
		code, err = encodeJump(s, mnem, pc+len(prefix), offsets, long, resolve)
	} else {
		code, ps, err = encodeNormal(s, fi, link)
	}
	if err != nil {
		return nil, nil, err
	}
	// Anchor the patch fields at function-relative positions: off indexes the
	// disp32 field, after is the address just past the instruction.
	body := pc + len(prefix)
	for i := range ps {
		ps[i].off += body
		ps[i].after = body + len(code)
	}
	return append(prefix, code...), ps, nil
}

func encodeNormal(s *ast.Instr, fi frameInfo, link *linkInfo) ([]byte, []sbPatch, error) {
	_, size := splitSize(strings.ToUpper(s.Mnemonic.Text))
	if size == 0 {
		size = 8
	}
	ops := make([]Operand, len(s.Operands))
	for i, op := range s.Operands {
		o, err := operandFromAST(op, size, fi, link)
		if err != nil {
			return nil, nil, err
		}
		ops[i] = o
	}
	e := &enc{}
	if err := e.encode(s.Mnemonic.Text, ops); err != nil {
		return nil, nil, err
	}
	ps := make([]sbPatch, len(e.patches))
	for i, p := range e.patches {
		ps[i] = sbPatch{off: p.off, name: p.name, addend: p.addend}
	}
	return e.out, ps, nil
}

// encodeJump encodes a JMP/CALL/Jcc with a relative offset resolved from the
// target label, in the short (rel8) or long (rel32) form.
func encodeJump(s *ast.Instr, mnem string, pc int, offsets map[string]int, long bool, resolve func(string) string) ([]byte, error) {
	if len(s.Operands) != 1 {
		return nil, fmt.Errorf("jump expects 1 operand, got %d", len(s.Operands))
	}
	name, ok := labelName(s.Operands[0])
	if !ok {
		return nil, fmt.Errorf("jump target must be a local label")
	}
	if resolve != nil && mnem != "CALL" {
		name = resolve(name)
	}
	target, ok := offsets[name]
	if !ok {
		return nil, fmt.Errorf("undefined label %q", name)
	}
	rel := int64(target - (pc + jumpSize(mnem, long)))

	if !long {
		if !fits8(rel) {
			return nil, fmt.Errorf("jump to %q does not fit the short form", name)
		}
		if mnem == "JMP" {
			return []byte{0xEB, byte(int8(rel))}, nil
		}
		cc, _ := condCode(mnem)
		return []byte{0x70 + byte(cc), byte(int8(rel))}, nil
	}
	switch mnem {
	case "JMP":
		return append([]byte{0xE9}, le32(rel)...), nil
	case "CALL":
		return append([]byte{0xE8}, le32(rel)...), nil
	default:
		cc, _ := condCode(mnem)
		return append([]byte{0x0F, 0x80 + byte(cc)}, le32(rel)...), nil
	}
}

// labelName extracts a local-label name from a jump operand.
func labelName(op *ast.Operand) (string, bool) {
	if op.Kind == ast.OpAddr && op.Addr.Sym != nil && op.Addr.Sym.Pseudo == "" &&
		op.Addr.Base == "" && op.Addr.Sym.Name != "" {
		return op.Addr.Sym.Name, true
	}
	return "", false
}

// spReg is the hardware stack pointer used to realise FP/SP pseudo-operands.
var spReg = Reg{idx: 4, size: 8}

// operandFromAST converts a parsed operand into an encoder Operand, applying
// the frame translation to FP/SP pseudo-register operands.
func operandFromAST(op *ast.Operand, size int, fi frameInfo, link *linkInfo) (Operand, error) {
	switch op.Kind {
	case ast.OpImmediate:
		if op.Imm.HasVal {
			v := op.Imm.Val
			if op.Imm.Neg {
				v = -v
			}
			return Imm(v), nil
		}
		return nil, fmt.Errorf("non-integer immediate not supported")

	case ast.OpAddr:
		a := op.Addr

		// FP-relative: x+N(FP) → (N + fpAdjust)(SP).  The offset N lives in the
		// symbol, not the address displacement.
		if a.Sym != nil && a.Sym.Pseudo == "FP" {
			off := a.Sym.Offset + fi.fpAdjust
			return Mem{Base: spReg, Disp: off, HasBase: true, Size: size}, nil
		}
		// SP-relative local: x-N(SP) → (spAdjust + offset)(SP).
		if a.Sym != nil && a.Sym.Pseudo == "SP" && a.Base == "" {
			off := fi.spAdjust + a.Sym.Offset
			return Mem{Base: spReg, Disp: off, HasBase: true, Size: size}, nil
		}
		// SB (global symbol): a symbol defined in the same file (GLOBL) is
		// encoded RIP-relative and resolved by the file-level layout;
		// anything not defined here needs object-file emission.
		if a.Sym != nil && a.Sym.Pseudo == "SB" {
			if link == nil || link.symbols == nil {
				return nil, fmt.Errorf("symbol %q needs file-level assembly (AssembleFile)", a.Sym.Name)
			}
			if !link.symbols[a.Sym.Name] {
				if a.Sym.Static {
					return nil, fmt.Errorf("undefined symbol %q", a.Sym.Name)
				}
				return nil, fmt.Errorf("external symbol %q needs object-file emission", a.Sym.Name)
			}
			return sbMem{size: size, name: a.Sym.Name, addend: a.Sym.Offset}, nil
		}

		// Memory with a real base register: (base), off(base), (base)(index*scale).
		if a.Base != "" {
			base, ok := ParseReg(a.Base)
			if !ok {
				return nil, fmt.Errorf("unknown base register %q", a.Base)
			}
			m := Mem{Base: base, Disp: a.Offset, HasBase: true, Size: size}
			if a.Index != "" {
				idx, ok := ParseReg(a.Index)
				if !ok {
					return nil, fmt.Errorf("unknown index register %q", a.Index)
				}
				m.Index = idx
				m.Scale = a.Scale
				m.HasIndex = true
			}
			return m, nil
		}
		// Bare register.
		if a.Sym != nil && a.Sym.Pseudo == "" && a.Sym.Name != "" {
			if r, ok := ParseReg(a.Sym.Name); ok {
				return r, nil
			}
		}
		return nil, fmt.Errorf("operand form not yet supported")
	}
	return nil, fmt.Errorf("unsupported operand")
}