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

package asm

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

// RISC-V frame mapping: translates Go's FP/SP pseudo-register addressing
// into real RISC-V memory accesses.
//
// In Go's ABI0 (used by assembly functions), arguments are passed on the
// stack.  At function entry the return address sits at SP, so the frame
// pointer FP == SP+8 and the first argument is at FP+0 == SP+8.
//
// On RISC-V the hardware registers are:
//   SP = X2 (stack pointer)
//   FP = S0 = X8 (frame pointer, by convention)
//
// For NOSPLIT $0 functions the prologue is omitted and arguments are read
// directly from SP+8+offset.

// riscvFrameInfo holds the frame parameters computed from a TEXT directive.
type riscvFrameInfo struct {
	frameSize int  // the $framesize from TEXT
	argsSize  int  // the -argsize from TEXT
	noSplit   bool // the NOSPLIT flag
}

// riscvComputeFrame extracts frame information from a TEXT directive.
func riscvComputeFrame(t *ast.Text) riscvFrameInfo {
	fi := riscvFrameInfo{}
	fi.frameSize = frameSize(t)
	fi.argsSize = argsSize(t)
	for _, f := range t.Flags {
		if f == "NOSPLIT" {
			fi.noSplit = true
		}
	}
	return fi
}

// riscvPrologue returns the prologue bytes for a RISC-V function.
// For NOSPLIT $0 functions there is no prologue.  For functions with a
// frame, we emit: ADDI SP, SP, -framesize; SD S0, (framesize-8)(SP); ...
func riscvPrologue(fi riscvFrameInfo) []byte {
	if fi.noSplit && fi.frameSize == 0 {
		return nil // no prologue for NOSPLIT $0
	}
	var out []byte
	if fi.frameSize > 0 {
		// ADDI SP, SP, -framesize
		out = append(out, riscvITypeLE(0x13, 0x0, 2, 2, int32(-fi.frameSize))...)
		// Save the frame pointer (S0 = X8) at the top of the new frame.
		// SD S0, (framesize-8)(SP)
		out = append(out, riscvSTypeLE(0x23, 0x3, 2, 8, int32(fi.frameSize-8))...)
	}
	return out
}

// riscvEpilogue returns the epilogue bytes for a RISC-V function.
func riscvEpilogue(fi riscvFrameInfo) []byte {
	if fi.noSplit && fi.frameSize == 0 {
		return nil
	}
	var out []byte
	if fi.frameSize > 0 {
		// Restore the frame pointer: LD S0, (framesize-8)(SP)
		out = append(out, riscvITypeLE(0x03, 0x3, 8, 2, int32(fi.frameSize-8))...)
		// ADDI SP, SP, framesize
		out = append(out, riscvITypeLE(0x13, 0x0, 2, 2, int32(fi.frameSize))...)
	}
	return out
}

// riscvResolvePseudo translates a pseudo-register memory reference into a
// real base register and offset.  It handles name+offset(FP) and
// name+offset(SP).
//
// Returns the base register number and the adjusted offset.
func riscvResolvePseudo(sym *ast.Symbol, fi riscvFrameInfo) (base int, off int32) {
	if sym == nil {
		return -1, 0
	}
	offset := int32(sym.Offset)
	switch sym.Pseudo {
	case "FP":
		// FP == SP+8 for NOSPLIT $0; arguments are at SP+8+offset.
		if fi.noSplit && fi.frameSize == 0 {
			return 2, 8 + offset // SP + 8 + argOffset
		}
		// With a frame, FP points to the saved frame; args are at FP+offset.
		return 8, offset // S0 + argOffset
	case "SP":
		// SP-relative; the offset is from the current SP.
		return 2, offset
	case "SB":
		// Static data reference — needs a relocation (not yet supported).
		return -1, offset
	default:
		return -1, offset
	}
}

// riscvITypeLE encodes an I-type instruction and returns little-endian bytes.
func riscvITypeLE(opcode, funct3 uint32, rd, rs1 int, imm int32) []byte {
	word := (uint32(imm&0xFFF) << 20) | (uint32(rs1) << 15) |
		(funct3 << 12) | (uint32(rd) << 7) | opcode
	return []byte{byte(word), byte(word >> 8), byte(word >> 16), byte(word >> 24)}
}

// riscvSTypeLE encodes an S-type instruction and returns little-endian bytes.
func riscvSTypeLE(opcode, funct3 uint32, rs1, rs2 int, imm int32) []byte {
	immU := uint32(imm) & 0xFFF
	word := ((immU >> 5) << 25) | (uint32(rs2) << 20) | (uint32(rs1) << 15) |
		(funct3 << 12) | ((immU & 0x1F) << 7) | opcode
	return []byte{byte(word), byte(word >> 8), byte(word >> 16), byte(word >> 24)}
}