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

package asm

import (
	"strings"

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

// RISC-V frame mapping, matching the Go toolchain's riscv64 backend.
//
// Go's riscv64 functions have no hardware frame pointer: FP and SP are
// synthetic registers resolved against the hardware stack pointer (X2) and
// the frame size.  The return address lives in the link register (X1, RA/LR).
//
// The autosize is the real stack adjustment: the declared local frame plus
// the 8 bytes for the saved link register (the toolchain's FixedFrameSize).
// A leaf function with a zero frame gets no prologue at all.
//
// Prologue (autosize > 0), byte-identical to the toolchain:
//
//	MOV  LR, -autosize(SP)   // save LR below the new SP (traceback-safe)
//	ADDI $-autosize, SP, SP  // open the frame
//	MOV  LR, 0(SP)           // save LR again at SP (signal-safety)
//
// Epilogue (autosize > 0): MOV 0(SP), LR; ADDI $autosize, SP, SP; the RET's
// uncompressed JALR X0, 0(X1) follows.  The toolchain restores LR on every
// frame, leaf or not.

// riscvFrameInfo holds the frame layout derived from a TEXT directive.
type riscvFrameInfo struct {
	autosize int // the real SP adjustment (locals + saved LR)
}

// riscvComputeFrame derives the frame layout for a TEXT function.
func riscvComputeFrame(t *ast.Text) riscvFrameInfo {
	frame := frameSize(t)
	if frame != 0 || !riscvIsLeaf(t) {
		// FixedFrameSize = 8: space for the saved link register.  A
		// zero-frame non-leaf function still opens an 8-byte frame for LR.
		return riscvFrameInfo{autosize: frame + 8}
	}
	return riscvFrameInfo{}
}

// riscvIsLeaf reports whether a function contains no call instructions.
// CALL always links; JAL/JALR link only when their destination register is
// the link register (X1), matching cmd/internal/obj/riscv's containsCall.
func riscvIsLeaf(t *ast.Text) bool {
	for _, stmt := range t.Body {
		in, ok := stmt.(*ast.Instr)
		if !ok {
			continue
		}
		switch strings.ToUpper(in.Mnemonic.Text) {
		case "CALL":
			return false
		case "JAL":
			// JAL rd, target — a call only when rd is the link register.
			if len(in.Operands) >= 2 && regFromOperand(in.Operands[0]) == 1 {
				return false
			}
		case "JALR":
			// JALR rs1, rd — a call when rd is X1; JALR offset(rs1) always
			// links to X1.
			if len(in.Operands) == 1 {
				return false
			}
			if len(in.Operands) >= 2 && regFromOperand(in.Operands[1]) == 1 {
				return false
			}
		}
	}
	return true
}

// riscvPrologue returns the prologue bytes for a RISC-V function, matching
// the toolchain's compression: the SP adjustment compresses to C.ADDI when
// the immediate fits, and the second LR save compresses to C.SDSP.
func riscvPrologue(fi riscvFrameInfo) []byte {
	if fi.autosize == 0 {
		return nil
	}
	var out []byte
	// MOV LR, -autosize(SP) — SD X1, -autosize(X2).  The negative offset is
	// not compressible to C.SDSP (unsigned), so it stays 4 bytes.
	out = append(out, wordLE(riscvSType(riscvEnc{0x23, 0x3, 0x00}, 2, 1, int32(-fi.autosize)))...)
	// ADDI $-autosize, SP, SP — open the frame (C.ADDI when it fits).
	out = append(out, riscvSPAdjust(int32(-fi.autosize))...)
	// MOV LR, 0(SP) — SD X1, 0(X2) → C.SDSP X1, 0.
	c := rvcSSP(0x7, 1, 0)
	out = append(out, byte(c), byte(c>>8))
	return out
}

// riscvReturn returns the bytes for a RET: the epilogue (restore LR and
// deallocate the frame when present) followed by the uncompressed JALR X0,
// 0(X1) the toolchain emits for RET (it never compresses RET to C.JR).
func riscvReturn(fi riscvFrameInfo) []byte {
	var out []byte
	if fi.autosize != 0 {
		// MOV 0(SP), LR — LD X1, 0(X2) → C.LDSP X1, 0.
		c := rvcLSP(0x3, 1, 0)
		out = append(out, byte(c), byte(c>>8))
		// ADDI $autosize, SP, SP — close the frame (C.ADDI when it fits).
		out = append(out, riscvSPAdjust(int32(fi.autosize))...)
	}
	// JALR X0, 0(X1).
	return append(out, wordLE(riscvIType(riscvEnc{0x67, 0x0, 0x00}, 0, 1, 0))...)
}

// riscvSPAdjust emits an ADDI rd, imm, rd for the stack pointer (rd = rs1 =
// X2), compressed to C.ADDI16SP when the immediate is a nonzero 16-byte
// multiple, else C.ADDI when it fits 6-bit signed.
func riscvSPAdjust(imm int32) []byte {
	if imm != 0 && imm%16 == 0 && imm >= -512 && imm <= 511 {
		c := rvcADDI16SP(2, imm)
		return []byte{byte(c), byte(c >> 8)}
	}
	if riscvFitsCAddi(imm) {
		c := rvcCI(0x0, 2, uint32(imm)&0x3F)
		return []byte{byte(c), byte(c >> 8)}
	}
	return wordLE(riscvIType(riscvEnc{0x13, 0x0, 0x00}, 2, 2, imm))
}

// riscvFitsCAddi reports whether imm compresses to C.ADDI (a nonzero 6-bit
// signed immediate).
func riscvFitsCAddi(imm int32) bool {
	return imm != 0 && imm >= -32 && imm <= 31
}

// riscvPrologueSpadjPC returns the function-relative byte offset where the
// prologue has finished decrementing SP (the delta becomes autosize).
func riscvPrologueSpadjPC(fi riscvFrameInfo) int {
	if fi.autosize == 0 {
		return 0
	}
	// SD (4 bytes) + ADDI/C.ADDI (2 or 4 bytes).
	return 4 + riscvSPAdjustLen(int32(-fi.autosize))
}

// riscvReturnEpilogueLen returns the byte length of the RET's epilogue up to
// (but not including) the final JALR — the point where SP is restored.
func riscvReturnEpilogueLen(fi riscvFrameInfo) int {
	if fi.autosize == 0 {
		return 0
	}
	// C.LDSP (2 bytes) + ADDI/C.ADDI (2 or 4 bytes).
	return 2 + riscvSPAdjustLen(int32(fi.autosize))
}

func riscvSPAdjustLen(imm int32) int {
	if imm != 0 && imm%16 == 0 && imm >= -512 && imm <= 511 {
		return 2
	}
	if riscvFitsCAddi(imm) {
		return 2
	}
	return 4
}

// riscvResolvePseudo translates a pseudo-register memory reference into a
// hardware base register and offset.  x+N(FP) → (N + autosize + 8)(SP);
// x+N(SP) → (N + autosize)(SP).  Returns base = -1 for an unresolvable
// reference (SB: static data, handled by the relocation path).
func riscvResolvePseudo(sym *ast.Symbol, fi riscvFrameInfo) (base int, off int32) {
	if sym == nil {
		return -1, 0
	}
	switch sym.Pseudo {
	case "FP":
		return 2, int32(sym.Offset) + int32(fi.autosize) + 8
	case "SP":
		return 2, int32(fi.autosize) + int32(sym.Offset)
	case "SB":
		return -1, int32(sym.Offset)
	}
	return -1, 0
}