// Copyright 2025 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package asm provides a simple parser for Go assembly files.
package asm

import (
	"bufio"
	"bytes"
	"fmt"
	"sort"
	"strings"
	"unicode"

	"golang.org/x/tools/gopls/internal/protocol"
)

// Kind describes the nature of an identifier in an assembly file.
type Kind uint8

const (
	Invalid Kind = iota // reserved zero value; not used by Ident
	Ref                 // arbitrary reference to symbol or control label
	Text                // definition of TEXT (function) symbol
	Global              // definition of GLOBL (var) symbol
	Data                // initialization of GLOBL (var) symbol; effectively a reference
	Label               // definition of control label
)

func (k Kind) String() string {
	if int(k) < len(kindString) {
		return kindString[k]
	}
	return fmt.Sprintf("Kind(%d)", k)
}

var kindString = [...]string{
	Invalid: "invalid",
	Ref:     "ref",
	Text:    "text",
	Global:  "global",
	Data:    "data",
	Label:   "label",
}

// A file represents a parsed file of Go assembly language.
type File struct {
	// Idents holds the identifiers of the file, ordered by Offset;
	// [File.IdentAt] relies on that order for its binary search.
	Idents []Ident

	Mapper *protocol.Mapper

	// TODO(adonovan): use token.File? This may be important in a
	// future in which analyzers can report diagnostics in .s files.
}

// IdentRange returns the protocol.Range for the identifier in this file.
func (f *File) IdentRange(ident Ident) (protocol.Range, error) {
	return f.Mapper.OffsetRange(ident.Offset, ident.Offset+ident.OrigLen)
}

// IdentLocation returns a protocol Location for the identifier in this file.
func (f *File) IdentLocation(ident Ident) (protocol.Location, error) {
	return f.Mapper.OffsetLocation(ident.Offset, ident.Offset+ident.OrigLen)
}

// IdentAt returns the identifier containing the byte range [start, end),
// or nil if none. Because [File.Idents] are ordered by Offset, the
// lookup uses a binary search.
func (f *File) IdentAt(start, end int) *Ident {
	// Find the last identifier whose Offset <= start.
	idx := sort.Search(len(f.Idents), func(i int) bool {
		return f.Idents[i].Offset > start
	})
	if idx == 0 {
		return nil
	}
	id := &f.Idents[idx-1]
	if end <= id.End() {
		return id
	}
	return nil
}

// FunctionRange returns the byte range [start, end) of the TEXT function
// enclosing offset: start is the beginning of the line containing the
// enclosing TEXT directive, end is the beginning of the line containing
// the next TEXT directive, or len(content) if there is none. If offset
// precedes the first TEXT directive, the range covers from 0 to the
// first TEXT directive.
//
// TEXT directives are taken from the parsed file rather than re-detected
// here, so that scoping stays consistent with the identifiers the parser
// reports (e.g. a bare "TEXT" line with no symbol is not a boundary).
func (f *File) FunctionRange(offset int) (int, int) {
	content := f.Mapper.Content
	funcStart, funcEnd := 0, len(content)
	for i := range f.Idents {
		id := &f.Idents[i]
		if id.Kind != Text {
			continue
		}
		lineStart := id.Offset
		for lineStart > 0 && content[lineStart-1] != '\n' {
			lineStart--
		}
		if lineStart > offset {
			funcEnd = lineStart
			break
		}
		funcStart = lineStart
	}
	return funcStart, funcEnd
}

// Ident represents an identifier in an assembly file.
type Ident struct {
	Name    string // symbol name (after correcting [·∕]); Name[0]='.' => current package
	Offset  int    // zero-based byte offset
	OrigLen int    // original length of the symbol name (before cleanup)
	Kind    Kind
}

// End returns the identifier's end offset.
func (id Ident) End() int { return id.Offset + id.OrigLen }

// Parse extracts identifiers from Go assembly files.
// Since it is a best-effort parser, it never returns an error.
func Parse(uri protocol.DocumentURI, content []byte) *File {
	var idents []Ident
	offset := 0 // byte offset of start of current line

	// TODO(adonovan) use a proper tokenizer that respects
	// comments, string literals, line continuations, etc.
	scan := bufio.NewScanner(bytes.NewReader(content))
	for ; scan.Scan(); offset += len(scan.Bytes()) + len("\n") {
		line := scan.Text()

		// Strip comments.
		if idx := strings.Index(line, "//"); idx >= 0 {
			line = line[:idx]
		}

		// Skip blank lines.
		if strings.TrimSpace(line) == "" {
			continue
		}

		// Check for label definitions (ending with colon).
		if colon := strings.IndexByte(line, ':'); colon > 0 {
			label := strings.TrimSpace(line[:colon])
			if isIdent(label) {
				idents = append(idents, Ident{
					Name:    label,
					Offset:  offset + strings.Index(line, label),
					OrigLen: len(label),
					Kind:    Label,
				})
				continue
			}
		}

		// Split line into words.
		words := strings.Fields(line)
		if len(words) == 0 {
			continue
		}

		// A line of the form
		//    TEXT ·sym<ABIInternal>(SB),NOSPLIT,$12
		// declares a text symbol "·sym".
		if len(words) > 1 {
			kind := Invalid
			switch words[0] {
			case "TEXT":
				kind = Text
			case "GLOBL":
				kind = Global
			case "DATA":
				kind = Data
			}
			if kind != Invalid {
				sym := words[1]
				sym = cutBefore(sym, ",") // strip ",NOSPLIT,$12" etc
				sym = cutBefore(sym, "(") // "sym(SB)" -> "sym"
				sym = cutBefore(sym, "<") // "sym<ABIInternal>" -> "sym"
				sym = strings.TrimSpace(sym)
				if isIdent(sym) {
					// (The Index call assumes sym is not itself "TEXT" etc.)
					idents = append(idents, Ident{
						Name:    cleanup(sym),
						Kind:    kind,
						Offset:  offset + strings.Index(line, sym),
						OrigLen: len(sym),
					})
				}
				continue
			}
		}

		// Find references in the rest of the line.
		pos := 0
		for _, word := range words {
			// Find actual position of word within line.
			tokenPos := strings.Index(line[pos:], word)
			if tokenPos < 0 {
				panic(line)
			}
			tokenPos += pos
			pos = tokenPos + len(word)

			// Reject probable instruction mnemonics (e.g. MOV).
			if len(word) >= 2 && word[0] != '·' &&
				!strings.ContainsFunc(word, unicode.IsLower) {
				continue
			}

			if word[0] == '$' {
				word = word[1:]
				tokenPos++

				// Reject probable immediate values (e.g. "$123").
				if !strings.ContainsFunc(word, isNonDigit) {
					continue
				}
			}

			// Reject probably registers (e.g. "PC").
			if len(word) <= 3 && !strings.ContainsFunc(word, unicode.IsLower) {
				continue
			}

			// Probable identifier reference.
			//
			// TODO(adonovan): handle FP symbols correctly;
			// sym+8(FP) is essentially a comment about
			// stack slot 8, not a reference to a symbol
			// with a declaration somewhere; so they form
			// an equivalence class without a canonical
			// declaration.
			//
			// TODO(adonovan): handle pseudoregisters and field
			// references such as:
			//    MOVD	$runtime·g0(SB), g      // pseudoreg
			//    MOVD	R0, g_stackguard0(g)    // field ref

			sym := cutBefore(word, "(") // "·sym(SB)" => "sym"
			sym = cutBefore(sym, "+")   // "sym+8(FP)" => "sym"
			sym = cutBefore(sym, "<")   // "sym<ABIInternal>" =>> "sym"
			if isIdent(sym) {
				idents = append(idents, Ident{
					Name:    cleanup(sym),
					Kind:    Ref,
					Offset:  offset + tokenPos,
					OrigLen: len(sym),
				})
			}
		}
	}

	_ = scan.Err() // ignore scan errors

	return &File{Idents: idents, Mapper: protocol.NewMapper(uri, content)}
}

// isIdent reports whether s is a valid Go assembly identifier.
func isIdent(s string) bool {
	for i, r := range s {
		if !isIdentRune(r, i) {
			return false
		}
	}
	return len(s) > 0
}

// cutBefore returns the portion of s before the first occurrence of sep, if any.
func cutBefore(s, sep string) string {
	if before, _, ok := strings.Cut(s, sep); ok {
		return before
	}
	return s
}

// cleanup converts a symbol name from assembler syntax to linker syntax.
func cleanup(sym string) string {
	return repl.Replace(sym)
}

var repl = strings.NewReplacer(
	"·", ".", // (U+00B7 MIDDLE DOT)
	"∕", "/", // (U+2215 DIVISION SLASH)
)

func isNonDigit(r rune) bool { return !unicode.IsDigit(r) }

// -- plundered from GOROOT/src/cmd/asm/internal/asm/parse.go --

// We want center dot (·) and division slash (∕) to work as identifier characters.
func isIdentRune(ch rune, i int) bool {
	if unicode.IsLetter(ch) {
		return true
	}
	switch ch {
	case '_': // Underscore; traditional.
		return true
	case '\u00B7': // Represents the period in runtime.exit. U+00B7 '·' middle dot
		return true
	case '\u2215': // Represents the slash in runtime/debug.setGCPercent. U+2215 '∕' division slash
		return true
	}
	// Digits are OK only after the first character.
	return i > 0 && unicode.IsDigit(ch)
}
