// 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, small), 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)
//
// Large frames (autosize past the 12-bit offset or immediate ranges) expand
// the store and the adjust through REGTMP (R30) exactly as the toolchain's
// assembler does: the store via the rounding LU12IW split, the adjust via
// the floor LU12IW/ORI split.
//
// Epilogue: MOVV 0(R3), R1; ADDV $autosize, R3 (non-leaf only for the LR
// restore; the adjust materialised when the immediate does not fit); 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; each materialisation shrinks by one word when the
// constant's low 12 bits are zero.
func loong64GuardLen(fi loong64FrameInfo) int {
	if !fi.needSplit {
		return 0
	}
	off := int64(fi.autosize - stackSmall)
	switch fi.splitClass {
	case 0:
		return 12
	case 1:
		if off <= 2048 {
			return 16 // ADDV $-off fits the signed 12-bit immediate
		}
		return 24 // MOVV + LU12IW + ORI + ADDV + SGTU + BEQ
	default:
		// MOVV + [mat] + SGTU + BNE + [mat] + ADDV + SGTU + BEQ
		return (6 + loong64MatLen(off) + loong64MatLen(-off)) * 4
	}
}

// loong64MatLen reports the word count of materialising v in R30: a value
// with a zero high part needs only the ORI (the toolchain's MOVW $v, R30),
// one with a zero low part only the LU12IW.
func loong64MatLen(v int64) int {
	if v>>12 == 0 || v&0xFFF == 0 {
		return 1
	}
	return 2
}

// loong64MatWords appends the words that materialise v in R30, splitting it
// as v>>12 plus the zero-extended low 12 bits.
func loong64MatWords(ws []uint32, v int64) []uint32 {
	hi := v >> 12
	lo := v & 0xFFF
	if hi == 0 {
		return append(ws, l64irr(l64OriOp, int(v), 0, 30))
	}
	ws = append(ws, l64ir(l64Lu12iwOp, int(hi), 30))
	if lo != 0 {
		ws = append(ws, l64irr(l64OriOp, int(lo), 30, 30))
	}
	return ws
}

// The LU12IW and ORI opcode bases (2RI20 and 2RI12 formats); the ORI reads
// and writes rd itself.
const (
	l64Lu12iwOp = 0x0a << 25
	l64OriOp    = 0x0e << 22
)

// loong64Imm12 reports whether v fits a signed 12-bit immediate.
func loong64Imm12(v int64) bool { return v >= -2048 && v <= 2047 }

// 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 and are computed from each
// branch's own position.
func loong64GuardBytes(fi loong64FrameInfo, blockStart int) []byte {
	// MOVV 16(g), R20 (g.stackguard0), g = R22.
	ws := []uint32{l64irr(l64loadStoreTable["MOVV"].ld, 16, 22, 20)}
	off := int64(fi.autosize - stackSmall)
	// beq appends BEQ R20, blockStart from the branch's own position.
	beq := func() {
		ws = append(ws, loong64Beqz(20, int32((blockStart-len(ws)*4)>>2)))
	}
	switch fi.splitClass {
	case 0:
		// SGTU SP, R20, R20; BEQ R20, more
		ws = append(ws, l64rrr(l64DualTable["SGTU"].rrr, 3, 20, 20))
		beq()
	case 1:
		ws = append(ws, loong64MediumWords(off)...)
		ws = append(ws, l64rrr(l64DualTable["SGTU"].rrr, 24, 20, 20))
		beq()
	default:
		// SGTU $off, SP, R24 catches the SP underflow a huge frame would
		// cause; BNE jumps to morestack in that case.
		ws = append(ws, loong64MatWords(nil, off)...)
		ws = append(ws, l64rrr(l64DualTable["SGTU"].rrr, 30, 3, 24))
		ws = append(ws, loong64Bnez(24, int32((blockStart-len(ws)*4)>>2)))
		ws = append(ws, loong64MatWords(nil, -off)...)
		ws = append(ws, l64rrr(l64DualTable["ADDV"].rrr, 30, 3, 24))
		ws = append(ws, l64rrr(l64DualTable["SGTU"].rrr, 24, 20, 20))
		beq()
	}
	return l64WordsLE(ws...)
}

// loong64MediumWords emits the medium-class stack check for offset off: the
// ADDV immediate when it fits, otherwise the same sequence with the constant
// materialised in R30.
func loong64MediumWords(off int64) []uint32 {
	if off <= 2048 {
		return []uint32{l64irr(l64DualTable["ADDV"].imm, int(-off), 3, 24)}
	}
	ws := loong64MatWords(nil, -off)
	return append(ws, l64rrr(l64DualTable["ADDV"].rrr, 30, 3, 24))
}

// 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, the
// toolchain's OR R1, R0, R31 expansion), BL runtime.morestack_noctxt, B back
// to the function entry.
func loong64MoreStackBlock(blockStart int) ([]byte, Reloc) {
	ws := []uint32{
		l64rrr(l64DualTable["OR"].rrr, 0, 1, 31), // MOVV R1, R31 (OR 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.  When
// the LR store offset leaves the toolchain's 12-bit store range ([-2046,
// 2045], BIG_12 = 2046) or the SP adjust immediate its 12-bit immediate
// range, each switches to the R30 materialisation the assembler expands it
// to: the store uses the rounding %hi/%lo split (LU12IW of (v+2048)>>12,
// REGTMP += SP, store at the raw offset), the adjust the floor split
// (LU12IW, ORI when the low part is non-zero, REGTMP += SP).
func loong64Prologue(fi loong64FrameInfo) []byte {
	if fi.autosize == 0 {
		return nil
	}
	addiD := l64DualTable["ADDV"].imm
	var ws []uint32
	storeBase := 3
	if fi.autosize > 2046 {
		// The store goes through REGTMP: LU12IW of the rounding split,
		// REGTMP += SP, then the store at REGTMP with the truncated offset.
		v := -int64(fi.autosize)
		ws = append(ws, l64ir(l64Lu12iwOp, int((v+2048)>>12), 30))
		ws = append(ws, l64rrr(l64DualTable["ADDV"].rrr, 3, 30, 30))
		storeBase = 30
	}
	ws = append(ws, l64irr(l64loadStoreTable["MOVV"].st, -fi.autosize, storeBase, 1)) // MOVV R1, -autosize(base)
	if loong64Imm12(-int64(fi.autosize)) {
		ws = append(ws, l64irr(addiD, -fi.autosize, 3, 3)) // ADDV $-autosize, R3
	} else {
		ws = append(ws, loong64MatWords(nil, -int64(fi.autosize))...)
		ws = append(ws, l64rrr(l64DualTable["ADDV"].rrr, 30, 3, 3))
	}
	ws = append(ws, l64irr(l64loadStoreTable["MOVV"].st, 0, 3, 1)) // MOVV R1, 0(R3)
	return l64WordsLE(ws...)
}

// 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 (materialised when the
		// immediate does not fit).
		if loong64Imm12(int64(fi.autosize)) {
			ws = append(ws, l64irr(l64DualTable["ADDV"].imm, fi.autosize, 3, 3))
		} else {
			ws = append(ws, loong64MatWords(nil, int64(fi.autosize))...)
			ws = append(ws, l64rrr(l64DualTable["ADDV"].rrr, 30, 3, 3))
		}
	}
	// jirl r0, r1, 0, return.
	ws = append(ws, l64irr16(l64branchTable["JIRL"], 0, 1, 0))
	return l64WordsLE(ws...)
}

// loong64StoreWords reports the prologue word count of the LR store, and
// loong64AdjustWords the word count of an SP adjust of v: the immediate
// forms when they fit, otherwise the R30 materialisation sequences.
func loong64StoreWords(autosize int) int {
	if autosize > 2046 {
		return 3
	}
	return 1
}

func loong64AdjustWords(v int64) int {
	if loong64Imm12(v) {
		return 1
	}
	return loong64MatLen(v) + 1
}

// loong64EpilogueWords reports the epilogue word count the RET expands to.
func loong64EpilogueWords(fi loong64FrameInfo) int {
	n := loong64AdjustWords(int64(fi.autosize))
	if !fi.leaf {
		n++
	}
	return n
}

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