// 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 always use the 32-bit
// relative form so instruction sizes are fixed and offsets resolve in a single
// layout pass.
//
// 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) {
	fi := computeFrame(t)

	// Pass 1: lay out instructions (including prologue/epilogue) to fix label
	// offsets.
	offsets := map[string]int{}
	sizes := make([]int, len(t.Body))
	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)
			if err != nil {
				return nil, nil, fmt.Errorf("%s: %w", s.Mnemonic.Text, err)
			}
			sizes[i] = sz
			pos += sz
		}
	}

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

// 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 (pass 1).  encodeInstr
// already includes the epilogue for a RET in a frame-pointer function; jumps use
// a fixed rel32 size (no epilogue).
func instrSize(s *ast.Instr, fi frameInfo) (int, error) {
	mnem := strings.ToUpper(s.Mnemonic.Text)
	if isJumpMnemonic(mnem) {
		return jumpSize(mnem), nil
	}
	code, err := encodeInstr(s, 0, nil, fi)
	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 fixed length of a rel32 jump instruction.
func jumpSize(mnem string) int {
	if mnem == "JMP" || mnem == "CALL" {
		return 5 // opcode + 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.
func encodeInstr(s *ast.Instr, pc int, offsets map[string]int, fi frameInfo) ([]byte, error) {
	mnem := strings.ToUpper(s.Mnemonic.Text)

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

	var code []byte
	var err error
	if isJumpMnemonic(mnem) {
		code, err = encodeJump(s, mnem, pc+len(prefix), offsets)
	} else {
		code, err = encodeNormal(s, fi)
	}
	if err != nil {
		return nil, err
	}
	return append(prefix, code...), nil
}

func encodeNormal(s *ast.Instr, fi frameInfo) ([]byte, 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)
		if err != nil {
			return nil, err
		}
		ops[i] = o
	}
	return Encode(s.Mnemonic.Text, ops...)
}

// encodeJump encodes a JMP/CALL/Jcc with a rel32 offset resolved from the
// target label.
func encodeJump(s *ast.Instr, mnem string, pc int, offsets map[string]int) ([]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")
	}
	target, ok := offsets[name]
	if !ok {
		return nil, fmt.Errorf("undefined label %q", name)
	}
	rel := int64(target - (pc + jumpSize(mnem)))

	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) (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) needs a relocation — not yet supported.
		if a.Sym != nil && a.Sym.Pseudo == "SB" {
			return nil, fmt.Errorf("SB (global symbol) operands need relocation support (pending)")
		}

		// 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")
}