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

// Loong64 frame mapping, matching the Go toolchain's loong64 backend.
//
// Go's loong64 functions have no frame pointer: FP and SP are synthetic
// registers resolved against the hardware stack pointer (R3) and the frame
// size.  The return address lives in R1 (the link register).
//
// The autosize is the real stack adjustment: the declared local frame plus
// the 8 bytes for the saved link register, rounded up to a multiple of 8
// (the toolchain aligns frames with `if autosize&4 != 0 { autosize += 4 }`).
// A leaf function (no calls) with a zero frame gets no prologue at all.
//
// Prologue (autosize > 0), byte-identical to the toolchain:
//
//	MOVV R1, -autosize(R3)   // save LR below the new SP (traceback-safe)
//	ADDV $-autosize, R3      // open the frame
//	MOVV R1, 0(R3)           // save LR again at SP (signal-safety)
//
// Epilogue: MOVV 0(R3), R1; ADDV $autosize, R3 (non-leaf only for the LR
// restore); the RET's jirl r0, r1, 0 follows.

// loong64FrameInfo holds the frame layout derived from a TEXT directive.
type loong64FrameInfo struct {
	autosize int  // the real SP adjustment (locals + saved LR, aligned)
	frame    int  // the declared $framesize
	args     int  // the declared -argsize
	noSplit  bool // the NOSPLIT flag
	leaf     bool // no call instructions in the body

	// Stack-split guard state: like amd64 and arm64, a leaf function with a
	// small autosize is auto-marked NOSPLIT by the toolchain.
	needSplit  bool
	splitClass int // 0: <=StackSmall, 1: <=StackBig, 2: >StackBig
}

// loong64ComputeFrame derives the frame layout for a TEXT function.
func loong64ComputeFrame(t *ast.Text) loong64FrameInfo {
	fi := loong64FrameInfo{
		frame: frameSize(t),
		args:  argsSize(t),
	}
	for _, f := range t.Flags {
		if f == "NOSPLIT" {
			fi.noSplit = true
		}
	}
	fi.leaf = loong64IsLeaf(t)
	if fi.frame != 0 {
		fi.autosize = fi.frame + 8 // space for the saved LR
		if fi.autosize&4 != 0 {
			fi.autosize += 4
		}
	} else if !fi.leaf {
		// A zero-frame non-leaf function still opens an 8-byte frame for LR.
		fi.autosize = 8
	}
	switch {
	case fi.noSplit:
	case fi.autosize < stackSmall && fi.leaf:
		// Auto-NOSPLIT, as the toolchain's leaf mark concludes.
	default:
		fi.needSplit = true
		switch {
		case fi.autosize <= stackSmall:
			fi.splitClass = 0
		case fi.autosize <= stackBig:
			fi.splitClass = 1
		default:
			fi.splitClass = 2
		}
	}
	return fi
}

// loong64GuardLen returns the byte length of the stack-split guard prefix
// (zero when the function needs no guard).  The big class materialises two
// constants through R30.
func loong64GuardLen(fi loong64FrameInfo) int {
	if !fi.needSplit {
		return 0
	}
	switch fi.splitClass {
	case 0:
		return 12
	case 1:
		return 16
	default:
		return 40 // MOVV + [LU12IW+ORI] + SGTU + BNE + [LU12IW+ORI] + ADDV + SGTU + BEQ
	}
}

// loong64Lu12iOri materialises the 32-bit constant v in rd with the
// toolchain's LU12IW/ORI pair (the ORI reads and writes rd itself).
func loong64Lu12iOri(rd int, v int64) []uint32 {
	hi := int32(v >> 12)
	lo := int32(v & 0xFFF)
	return []uint32{
		0x0a<<25 | uint32(hi&0xFFFFF)<<5 | uint32(rd),
		0x0e<<22 | uint32(lo)<<10 | uint32(rd)<<5 | uint32(rd),
	}
}

// loong64GuardBytes emits the stack-split guard prefix.  blockStart is the
// function-relative address of the morestack call at the end of the function;
// branch displacements are in instructions.
func loong64GuardBytes(fi loong64FrameInfo, blockStart int) []byte {
	// MOVV 16(g), R20 (g.stackguard0), g = R22.
	ws := []uint32{l64irr(l64loadStoreTable["MOVV"].ld, 16, 22, 20)}
	switch fi.splitClass {
	case 0:
		// SGTU SP, R20, R20; BEQ R20, more
		ws = append(ws, l64rrr(l64DualTable["SGTU"].rrr, 3, 20, 20))
		ws = append(ws, loong64Beqz(20, int32((blockStart-8)>>2)))
	case 1:
		off := int32(fi.autosize - stackSmall)
		ws = append(ws, l64irr(l64DualTable["ADDV"].imm, int(-off), 3, 24))
		ws = append(ws, l64rrr(l64DualTable["SGTU"].rrr, 24, 20, 20))
		ws = append(ws, loong64Beqz(20, int32((blockStart-12)>>2)))
	default:
		off := int64(fi.autosize - stackSmall)
		movLen := 8 // LU12IW + ORI
		ws = append(ws, loong64Lu12iOri(30, off)...)
		ws = append(ws, l64rrr(l64DualTable["SGTU"].rrr, 30, 3, 24))
		ws = append(ws, loong64Bnez(24, int32((blockStart-(8+movLen))>>2)))
		ws = append(ws, loong64Lu12iOri(30, -off)...)
		ws = append(ws, l64rrr(l64DualTable["ADDV"].rrr, 30, 3, 24))
		ws = append(ws, l64rrr(l64DualTable["SGTU"].rrr, 24, 20, 20))
		ws = append(ws, loong64Beqz(20, int32((blockStart-loong64GuardLen(fi)+12)>>2)))
	}
	return l64WordsLE(ws...)
}

// loong64Beqz/loong64Bnez build the 21-bit conditional branches against R0
// that the toolchain emits for its guard compares.
func loong64Beqz(rj int, dispInstr int32) uint32 {
	return l64ir21(l64branch21Table["BEQZ"], int(dispInstr), rj)
}

func loong64Bnez(rj int, dispInstr int32) uint32 {
	return l64ir21(l64branch21Table["BNEZ"], int(dispInstr), rj)
}

// loong64MoreStackBlock emits the trailing block: MOVV R1, R31 (save LR),
// BL runtime.morestack_noctxt, B back to the function entry.
func loong64MoreStackBlock(blockStart int) ([]byte, Reloc) {
	ws := []uint32{
		l64rrr(l64DualTable["ADD"].rrr, 0, 1, 31), // MOVV R1, R31 (ADD R1, R0, R31)
		l64bbl(l64jumpTable["BL"], 0),             // BL, patched by the linker
	}
	disp := (-(blockStart + 8)) >> 2
	ws = append(ws, l64bbl(l64jumpTable["B"], int(disp)))
	reloc := Reloc{
		Off:   blockStart + 4,
		After: blockStart + 8,
		Name:  "runtime\u00b7morestack_noctxt",
		Kind:  RelLoong64Branch,
	}
	return l64WordsLE(ws...), reloc
}

// loong64IsLeaf reports whether a function contains no call instructions
// (JAL/BL/CALL), matching the toolchain's LEAF mark, which drives the frame
// and the epilogue shape.
func loong64IsLeaf(t *ast.Text) bool {
	for _, stmt := range t.Body {
		in, ok := stmt.(*ast.Instr)
		if !ok {
			continue
		}
		switch strings.ToUpper(in.Mnemonic.Text) {
		case "JAL", "CALL", "BL":
			return false
		}
	}
	return true
}

// loong64Prologue returns the prologue bytes for a loong64 function.
func loong64Prologue(fi loong64FrameInfo) []byte {
	if fi.autosize == 0 {
		return nil
	}
	addiD := l64DualTable["ADDV"].imm
	return l64WordsLE(
		l64irr(l64loadStoreTable["MOVV"].st, -fi.autosize, 3, 1), // MOVV R1, -autosize(R3)
		l64irr(addiD, -fi.autosize, 3, 3),                        // ADDV $-autosize, R3
		l64irr(l64loadStoreTable["MOVV"].st, 0, 3, 1),            // MOVV R1, 0(R3)
	)
}

// loong64Return returns the bytes for a RET: the epilogue (restore LR and
// deallocate the frame when present) followed by jirl r0, r1, 0.
func loong64Return(fi loong64FrameInfo) []byte {
	var ws []uint32
	if fi.autosize != 0 {
		if !fi.leaf {
			// MOVV 0(R3), R1, restore the link register.
			ws = append(ws, l64irr(l64loadStoreTable["MOVV"].ld, 0, 3, 1))
		}
		// ADDV $autosize, R3, close the frame.
		ws = append(ws, l64irr(l64DualTable["ADDV"].imm, fi.autosize, 3, 3))
	}
	// jirl r0, r1, 0, return.
	ws = append(ws, l64irr16(l64branchTable["JIRL"], 0, 1, 0))
	return l64WordsLE(ws...)
}

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