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

// The language-server side of the extended-instruction layer.  The registry
// in asm carries the mnemonics the generated architecture tables do not know;
// this file lets completion and hover document those mnemonics from the
// layer's own metadata, and lets diagnostics leave the layer's verdicts to
// it.  Nothing the server states about a registered mnemonic is invented:
// the summaries, forms, features, encodings and references are read from the
// metadata, and the operand spellings are measured against the layer's own
// encoder, so the editor surface can never promise a spelling the assembler
// would refuse.

package lsp

import (
	"fmt"
	"slices"
	"strings"
	"sync"

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

// extFormDoc is the rendering of one registered form: everything hover and
// completion state about it.
type extFormDoc struct {
	summary string // the metadata's one-line description ("" when absent)
	form    string // the operand shape's label, from ExtForm.String
	shape   string // the operand spellings the layer's encoder accepts
	feature string // the architecture feature, from ExtFeature
	word    uint32 // the fixed encoding bits
	ref     string // the manual entry the encoding is transcribed from
}

// extDocsKey addresses one mnemonic's renderings for one architecture.
type extDocsKey struct {
	a    arch.Arch
	name string
}

// extDocsCache memoises the per-form renderings: the probe walks the layer's
// encoder once per mnemonic and architecture, and every later completion or
// hover reuses the answer.
var extDocsCache sync.Map // map[extDocsKey][]extFormDoc

// The probe candidates.  They cover every operand kind the layer defines
// today; a kind arriving later joins through its case in extBuild, and an
// unknown kind degrades to a bare label rather than a wrong spelling.
var extProbeArrangements = []arch.ExtArrangement{
	arch.ExtArrB, arch.ExtArrH, arch.ExtArrS, arch.ExtArrD, arch.ExtArrQ, arch.ExtArrNone,
}

var extProbeQualifiers = []arch.ExtQualifier{
	arch.ExtQualMerging, arch.ExtQualZeroing, arch.ExtQualNone,
}

var extProbeImmediates = []int64{1, 0, 255, -1}

// extForms returns the renderings of one mnemonic's registered forms on a,
// computing and memoising them on first use.  It returns nil when the layer
// registers nothing for the mnemonic.
func extForms(a arch.Arch, mnemonic string) []extFormDoc {
	key := extDocsKey{a: a, name: strings.ToUpper(mnemonic)}
	if v, ok := extDocsCache.Load(key); ok {
		docs, _ := v.([]extFormDoc)
		return docs
	}
	var docs []extFormDoc
	if cands, ok := asm.LookupExtension(a, mnemonic); ok {
		docs = make([]extFormDoc, 0, len(cands))
		for _, in := range cands {
			docs = append(docs, extFormDoc{
				summary: in.Summary,
				form:    in.Form.String(),
				shape:   extFormShape(in),
				feature: string(in.Feature),
				word:    in.Word,
				ref:     in.Ref,
			})
		}
	}
	extDocsCache.Store(key, docs)
	return docs
}

// extBuild builds one candidate operand list for the form: every vector
// position at reg with the arrangement arr, every predicate position at reg
// with the qualifier qual, every immediate at imm.
func extBuild(in arch.ExtInstr, reg int, arr arch.ExtArrangement, qual arch.ExtQualifier, imm int64) []arch.ExtOperand {
	kinds := in.Form.Kinds()
	ops := make([]arch.ExtOperand, len(kinds))
	for i, k := range kinds {
		switch k {
		case arch.ExtZReg:
			ops[i] = arch.ExtVector(reg, arr)
		case arch.ExtPReg:
			ops[i] = arch.ExtPredicate(reg, qual)
		default: // arch.ExtImm and anything the probe does not model
			ops[i] = arch.ExtImmediate(imm)
		}
	}
	return ops
}

// extEncodes reports whether the form's encoder accepts the operand list.
func extEncodes(in arch.ExtInstr, ops []arch.ExtOperand) bool {
	_, err := in.Encode(ops)
	return err == nil
}

// extAnchor finds one operand list the form encodes, to anchor the probes on;
// ok is false when no probe combination encodes, and the shape then renders
// from the operand kinds alone.
func extAnchor(in arch.ExtInstr) (base []arch.ExtOperand, arr arch.ExtArrangement, qual arch.ExtQualifier, imm int64, ok bool) {
	for _, a := range extProbeArrangements {
		for _, q := range extProbeQualifiers {
			for _, v := range extProbeImmediates {
				ops := extBuild(in, 0, a, q, v)
				if extEncodes(in, ops) {
					return ops, a, q, v, true
				}
			}
		}
	}
	return nil, arch.ExtArrNone, arch.ExtQualNone, 0, false
}

// extAcceptedArrangements returns the arrangements the form's vector
// positions accept, probed together: the encoder requires the vector
// positions of a form to agree, so one position alone cannot carry the probe.
func extAcceptedArrangements(in arch.ExtInstr, qual arch.ExtQualifier, imm int64) []arch.ExtArrangement {
	var out []arch.ExtArrangement
	for _, a := range extProbeArrangements {
		if extEncodes(in, extBuild(in, 0, a, qual, imm)) {
			out = append(out, a)
		}
	}
	return out
}

// extAcceptedQualifiers returns the qualifiers the form's predicate
// positions accept, probed together for the same reason.
func extAcceptedQualifiers(in arch.ExtInstr, arr arch.ExtArrangement, imm int64) []arch.ExtQualifier {
	var out []arch.ExtQualifier
	for _, q := range extProbeQualifiers {
		if extEncodes(in, extBuild(in, 0, arr, q, imm)) {
			out = append(out, q)
		}
	}
	return out
}

// extRegRange walks the register numbers one operand position accepts and
// returns a "Z0-Z31" style span when the accepted set is one contiguous run
// from its low end.  ok is false when no register encodes there or the set
// is not contiguous, and the caller falls back to a bare label.
func extRegRange(in arch.ExtInstr, base []arch.ExtOperand, pos int, letter string) (span string, ok bool) {
	const limit = 32 // both register files the layer defines sit inside 32
	lo, hi, n := -1, -1, 0
	for r := range limit {
		ops := slices.Clone(base)
		ops[pos].Reg = r
		if !extEncodes(in, ops) {
			continue
		}
		n++
		if lo < 0 {
			lo = r
		}
		hi = r
	}
	if n == 0 || hi-lo+1 != n {
		return "", false
	}
	return fmt.Sprintf("%s%d-%s%d", letter, lo, letter, hi), true
}

// extFormShape measures the operand spellings the form accepts: the register
// spans, the arrangement suffixes and the predicate qualifiers its own
// encoder takes, rendered the way the source writes them.  Where a probe
// finds nothing to anchor on, the shape degrades to the operand kinds and
// states nothing the encoder has not proven.
func extFormShape(in arch.ExtInstr) string {
	kinds := in.Form.Kinds()
	if len(kinds) == 0 {
		return ""
	}
	base, arr, qual, imm, anchored := extAnchor(in)

	// The arrangement suffixes and the qualifiers, rendered in enum order.
	suffix := ""
	quals := []arch.ExtQualifier{qual}
	if anchored {
		var spelled []string
		for _, a := range extAcceptedArrangements(in, qual, imm) {
			if a != arch.ExtArrNone {
				spelled = append(spelled, a.String())
			}
		}
		if len(spelled) > 0 {
			suffix = strings.Join(spelled, "/")
		}
		quals = extAcceptedQualifiers(in, arr, imm)
	}

	var parts []string
	for i, k := range kinds {
		switch k {
		case arch.ExtZReg:
			label := "Z<n>"
			if anchored {
				if span, ok := extRegRange(in, base, i, "Z"); ok {
					label = span
				}
			}
			parts = append(parts, label+suffix)
		case arch.ExtPReg:
			label := "P<n>"
			if anchored {
				if span, ok := extRegRange(in, base, i, "P"); ok {
					label = span
				}
			}
			var spelled []string
			for _, q := range quals {
				if q != arch.ExtQualNone {
					spelled = append(spelled, q.String())
				}
			}
			if len(spelled) > 0 {
				label += strings.Join(spelled, " or ")
			}
			parts = append(parts, label)
		default:
			parts = append(parts, "$imm")
		}
	}
	return strings.Join(parts, ", ")
}

// extFeatureList names the architecture features the forms belong to, in
// first-occurrence order, "" when the metadata carries none.
func extFeatureList(docs []extFormDoc) string {
	var out []string
	for _, d := range docs {
		if d.feature != "" && !slices.Contains(out, d.feature) {
			out = append(out, d.feature)
		}
	}
	return strings.Join(out, ", ")
}

// renderExtensionHover renders the hover documentation of a registered
// extended mnemonic: the fact that it is an extension above the toolchain,
// then one entry per registered form.  A form with no summary in the
// metadata states the encoding facts alone: the shape label, the feature,
// the fixed encoding bits and the manual reference, and nothing more.
func renderExtensionHover(mnemonic string, docs []extFormDoc) string {
	if len(docs) == 0 {
		return ""
	}
	var b strings.Builder
	b.WriteString("**" + mnemonic + "**: extension above the Go toolchain")
	if f := extFeatureList(docs); f != "" {
		b.WriteString(" (" + f + ")")
	}
	for _, d := range docs {
		b.WriteString("\n\n- ")
		// The summary leads where the metadata carries one; without it the
		// form's own label stands in, and the bullet states the encoding
		// facts and nothing more.
		lead := d.summary
		if lead == "" {
			lead = d.form
		}
		if lead != "" {
			b.WriteString(lead + ": ")
		}
		b.WriteString(d.shape)
		var clauses []string
		if d.summary != "" && d.form != "" {
			clauses = append(clauses, d.form)
		}
		if d.feature != "" {
			clauses = append(clauses, d.feature)
		}
		if d.word != 0 {
			clauses = append(clauses, fmt.Sprintf("encoding 0x%08x", d.word))
		}
		if len(clauses) > 0 {
			b.WriteString(" (" + strings.Join(clauses, ", ") + ")")
		}
		if d.ref != "" {
			b.WriteString("\n  " + d.ref)
		}
	}
	return b.String()
}

// extensionHover returns the hover documentation of a registered extended
// mnemonic, "" when the layer registers nothing for it.
func extensionHover(a arch.Arch, word string) string {
	return renderExtensionHover(strings.ToUpper(word), extForms(a, word))
}

// applyExtensionCompletion merges the registered extended mnemonics into the
// completion list, so they are offered exactly where the toolchain
// mnemonics are: a mnemonic the base table also knows enriches its existing
// item, a registry-only one gains its own.  The detail carries the operand
// shapes the layer's encoder accepts.
func applyExtensionCompletion(items []CompletionItem, a arch.Arch) []CompletionItem {
	byLabel := make(map[string]int, len(items))
	for i, it := range items {
		if _, ok := byLabel[it.Label]; !ok {
			byLabel[it.Label] = i
		}
	}
	for _, name := range asm.ExtensionNames(a) {
		docs := extForms(a, name)
		if len(docs) == 0 {
			continue
		}
		shapes := make([]string, 0, len(docs))
		for _, d := range docs {
			shapes = append(shapes, d.shape)
		}
		hover := extensionHover(a, name)
		if i, ok := byLabel[name]; ok {
			items[i].Detail += "; extension: " + strings.Join(shapes, "; ")
			if items[i].Documentation == "" {
				items[i].Documentation = hover
			} else {
				items[i].Documentation += "\n\n" + hover
			}
			continue
		}
		byLabel[name] = len(items)
		items = append(items, CompletionItem{
			Label:         name,
			Kind:          ciFunction,
			Detail:        "extension: " + strings.Join(shapes, "; "),
			Documentation: hover,
		})
	}
	return items
}

// extOwnedCode reports whether a lint code carries a verdict the extension
// layer owns: the two that accuse a mnemonic of not existing or of not being
// encodable.  The operand-count rule stays with the generated table, whose
// bounds the relaxed architectures never fire against a registered form.
func extOwnedCode(code string) bool {
	return code == lint.CodeUnknownInstr || code == lint.CodeUnencodable
}

// extensionMnemonicSites collects the source positions of the statements
// whose mnemonic the architecture's extension layer registers.  The
// diagnostics pass consults it to leave the layer's verdicts to the layer:
// a registry mnemonic is neither unknown nor unencodable, the registry
// encodes it.
func extensionMnemonicSites(f *ast.File, a arch.Arch) map[[2]int]bool {
	out := make(map[[2]int]bool)
	if f == nil {
		return out
	}
	for _, d := range f.Decls {
		t, ok := d.(*ast.Text)
		if !ok {
			continue
		}
		for _, s := range t.Body {
			in, ok := s.(*ast.Instr)
			if !ok {
				continue
			}
			if _, ok := asm.LookupExtension(a, in.Mnemonic.Text); ok {
				out[[2]int{in.Mnemonic.Pos.Line, in.Mnemonic.Pos.Column}] = true
			}
		}
	}
	return out
}