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

// The assembler's side of the extended-instruction layer: this file turns a
// parsed arm64 statement into the operand form arch.ExtInstr.Encode consumes
// and routes statements only the layer can encode through the registry.  It
// sits beside the main arm64 encoders, never inside them: the generated
// tables and the scalar, NEON and FP paths are untouched, and a statement
// reaches this file only when the mnemonic is registered in the extension
// layer and at least one operand is a scalable vector or predicate register.
//
// The spellings are the layer's own Plan 9 forms, the ones its metadata
// documents: Zn, Zm, Zd for the unpredicated three-vector class, Zm, Pg/M,
// Zdn for the predicated class, imm{, LSL #8}, Zdn for the immediate classes.

package asm

import (
	"fmt"
	"strconv"
	"strings"

	"sourcedock.dev/petrbalvin/gasm-sdk/arch"
	"sourcedock.dev/petrbalvin/gasm-sdk/ast"
)

// arm64ExtStatement converts one instruction's operands into the extended
// layer's operand form.  pinned reports that the statement belongs to the
// layer: the mnemonic is registered in the registry and the operand list
// carries at least one scalable vector or predicate register.  A pinned
// statement can only encode through the layer, so every operand is read
// here and its diagnostic replaces whatever the scalar paths would have
// said about operands they cannot read; err is non-nil for a pinned
// statement whose operands the layer refuses, and extops is complete only
// when err is nil.  Unpinned means the statement is nobody's: the caller
// falls through to the ordinary arm64 encoders, which keep their exact
// behaviour for every scalar, NEON and FP operand list.
func arm64ExtStatement(mnem string, ops []*ast.Operand) (extops []arch.ExtOperand, pinned bool, err error) {
	if _, ok := LookupExtension(arch.ARM64, mnem); !ok {
		return nil, false, nil
	}
	if !arm64ExtPinned(ops) {
		return nil, false, nil
	}
	out := make([]arch.ExtOperand, 0, len(ops))
	for i, op := range ops {
		text := strings.Join(strings.Fields(op.Raw), "")

		// The spelled shift of an immediate class: the shift is an attribute
		// of the preceding immediate operand (imm{, LSL #8}, Zdn), never an
		// operand of its own.
		if amount, ok := strings.CutPrefix(text, "LSL#"); ok {
			if len(out) == 0 || out[len(out)-1].Kind != arch.ExtImm || out[len(out)-1].HasShift {
				return nil, true, fmt.Errorf("%s: operand %d (%s): LSL belongs straight after an immediate", mnem, i+1, op.Raw)
			}
			n, convErr := strconv.Atoi(amount)
			if convErr != nil {
				return nil, true, fmt.Errorf("%s: operand %d (%s): %q is not an LSL amount", mnem, i+1, op.Raw, amount)
			}
			out[len(out)-1].Shift, out[len(out)-1].HasShift = n, true
			continue
		}

		if op.Kind == ast.OpImmediate {
			ext, ok := arm64ExtImmediate(op)
			if !ok {
				return nil, true, fmt.Errorf("%s: operand %d (%s) is not an immediate the layer can read", mnem, i+1, op.Raw)
			}
			out = append(out, ext)
			continue
		}
		if ext, ok := arm64ExtVector(text); ok {
			out = append(out, ext)
			continue
		}
		if ext, ok := arm64ExtPredicate(text); ok {
			out = append(out, ext)
			continue
		}
		return nil, true, fmt.Errorf("%s: operand %d (%s) is not an extended-layer operand: want a scalable vector register, a predicate register or an immediate", mnem, i+1, op.Raw)
	}
	return out, true, nil
}

// arm64ExtPinned reports whether any operand is a scalable vector or
// predicate register, the shapes only the extension layer reads.  The test
// is deliberately loose about the suffixes: P0/B is not a spelling the
// layer takes, but the P of it makes the statement the layer's, and the
// conversion then diagnoses the operand precisely instead of leaving it to
// a scalar path that would report an unrelated register error.
func arm64ExtPinned(ops []*ast.Operand) bool {
	for _, op := range ops {
		if op.Kind == ast.OpImmediate {
			continue
		}
		text := strings.Join(strings.Fields(op.Raw), "")
		if _, ok := arm64ExtVector(text); ok {
			return true
		}
		if arm64ExtPredicateShape(text) {
			return true
		}
	}
	return false
}

// arm64ExtPredicateShape reports whether text spells a predicate register at
// all: P, digits, an optional arrangement suffix and an optional qualifier
// after a slash, whatever the qualifier says.  The strict parse in
// arm64ExtPredicate judges the suffix; this shape only decides who the
// operand belongs to.
func arm64ExtPredicateShape(text string) bool {
	if text == "" || text[0] != 'P' {
		return false
	}
	text = text[1:]
	if i := strings.IndexByte(text, '/'); i >= 0 {
		text = text[:i]
	}
	if i := strings.IndexByte(text, '.'); i >= 0 {
		text = text[:i]
	}
	_, err := strconv.Atoi(text)
	return err == nil && text != ""
}

// arm64ExtImmediate converts a $ immediate into the layer's form.  The
// parser folds a parenthesised constant expression in full ($(255<<8)) and
// reads a bare literal greedily, dropping any trailing operator tokens:
// $255<<8 parses as 255 with the shift silently gone.  Encoding that silent
// prefix would assemble what the text did not say, so an unparenthesised
// immediate is accepted only when its whole text reads back as one integer
// carrying the parser's value.
func arm64ExtImmediate(op *ast.Operand) (arch.ExtOperand, bool) {
	if op.Kind != ast.OpImmediate || !op.Imm.HasVal {
		return arch.ExtOperand{}, false
	}
	text := strings.Join(strings.Fields(strings.TrimPrefix(op.Raw, "$")), "")
	if !strings.HasPrefix(text, "(") {
		if _, parseErr := strconv.ParseInt(text, 0, 64); parseErr != nil {
			return arch.ExtOperand{}, false
		}
	}
	v := op.Imm.Val
	if op.Imm.Neg {
		v = -v
	}
	return arch.ExtOperand{Kind: arch.ExtImm, Imm: v}, true
}

// arm64ExtVector parses a scalable vector register operand: Z0..Z31 with an
// optional element-size suffix, Z0.S.  The arrangement is carried as written
// and the encoding validates it against the form.
func arm64ExtVector(text string) (arch.ExtOperand, bool) {
	reg, arr, ok := arm64ExtReg(text, 'Z')
	if !ok {
		return arch.ExtOperand{}, false
	}
	return arch.ExtOperand{Kind: arch.ExtZReg, Reg: reg, Arr: arr}, true
}

// arm64ExtPredicate parses a predicate register operand: P0..P15 with an
// optional element-size suffix and an optional qualifier, P0/M, P0.Z, P0.B/M.
func arm64ExtPredicate(text string) (arch.ExtOperand, bool) {
	qual := arch.ExtQualNone
	if base, suffix, found := strings.Cut(text, "/"); found {
		switch suffix {
		case "M":
			qual = arch.ExtQualMerging
		case "Z":
			qual = arch.ExtQualZeroing
		default:
			return arch.ExtOperand{}, false
		}
		text = base
	}
	reg, arr, ok := arm64ExtReg(text, 'P')
	if !ok {
		return arch.ExtOperand{}, false
	}
	return arch.ExtOperand{Kind: arch.ExtPReg, Reg: reg, Arr: arr, Qual: qual}, true
}

// arm64ExtReg parses Pn or Zn with an optional arrangement suffix off a
// normalised operand text.  The register range is left to the encoding: the
// layer's own diagnostics name the range a form carries.
func arm64ExtReg(text string, letter byte) (reg int, arr arch.ExtArrangement, ok bool) {
	if len(text) < 2 || text[0] != letter {
		return 0, 0, false
	}
	digits := text[1:]
	if base, suffix, found := strings.Cut(digits, "."); found {
		switch suffix {
		case "B":
			arr = arch.ExtArrB
		case "H":
			arr = arch.ExtArrH
		case "S":
			arr = arch.ExtArrS
		case "D":
			arr = arch.ExtArrD
		case "Q":
			arr = arch.ExtArrQ
		default:
			return 0, 0, false
		}
		digits = base
	}
	n, err := strconv.Atoi(digits)
	if err != nil || n < 0 {
		return 0, 0, false
	}
	return n, arr, true
}