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

// Command gasm is the developer frontend for GAsm — Go's Plan 9 assembler.
// It bundles a token dumper, a parser, a formatter, a linter and a language
// server into one binary.  Every subcommand works headlessly so it can be
// driven from scripts and CI as well as from an editor.
package main

import (
	"flag"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"strings"

	"sourcedock.dev/petrbalvin/gasm-devkit/arch"
	"sourcedock.dev/petrbalvin/gasm-devkit/asm"
	"sourcedock.dev/petrbalvin/gasm-devkit/ast"
	"sourcedock.dev/petrbalvin/gasm-devkit/format"
	"sourcedock.dev/petrbalvin/gasm-devkit/lexer"
	"sourcedock.dev/petrbalvin/gasm-devkit/lint"
	"sourcedock.dev/petrbalvin/gasm-devkit/lsp"
	"sourcedock.dev/petrbalvin/gasm-devkit/parser"
)

// version is the release version, stamped at build time via
// -ldflags "-X main.version=…" (defaulting to the current release).
var version = "0.13.0"

func main() {
	if len(os.Args) < 2 {
		usage(os.Stderr)
		os.Exit(2)
	}
	switch os.Args[1] {
	case "tokens":
		os.Exit(cmdTokens(os.Args[2:]))
	case "parse":
		os.Exit(cmdParse(os.Args[2:]))
	case "fmt":
		os.Exit(cmdFmt(os.Args[2:]))
	case "lint":
		os.Exit(cmdLint(os.Args[2:]))
	case "asm":
		os.Exit(cmdAsm(os.Args[2:]))
	case "lsp":
		os.Exit(cmdLSP(os.Args[2:]))
	case "version", "--version", "-V":
		os.Exit(cmdVersion())
	case "help", "--help", "-h":
		usage(os.Stdout)
	default:
		fmt.Fprintf(os.Stderr, "gasm: unknown command %q — run \"gasm --help\" for usage\n", os.Args[1])
		os.Exit(2)
	}
}

// cmdVersion prints the release version.
func cmdVersion() int {
	fmt.Printf("gasm %s\n", version)
	return 0
}

func usage(w io.Writer) {
	fmt.Fprintf(w, `gasm %s — developer tooling for Go's Plan 9 assembler (GAsm)

gasm bundles a lexer, parser, formatter, linter, standalone assembler and
language server for Plan 9 assembly into one self-contained binary.

Usage:
  gasm <command> [arguments]
  gasm [flags]

Commands:
  tokens    print the lexical token stream
  parse     parse and report syntax errors
  fmt       canonicalise formatting (gofmt for assembly)
  lint      run static checks
  asm       assemble .s files to machine code (amd64)
  lsp       run the language server over stdio
  version   print the version (same as --version)

Flags:
  -h, --help      show this help
  -V, --version   print the version

Run "gasm <command> -h" for a command's usage and flags.

Examples:
  gasm fmt                  reformat every .s below the current directory
  gasm lint go-flac/*.s     run static checks over the kernels
  gasm asm -o k.bin kern_amd64.s
  gasm asm --format elf -o k.o kern_amd64.s
  gasm asm --format goobj -p pkg/path -o k.o kern_amd64.s
`, version)
}

// newCommand returns the FlagSet of a subcommand whose -h/--help prints a
// proper usage block: the one-line usage, the long description and the flag
// defaults.  The flag package routes -h/--help to fs.Usage and exits 0.
func newCommand(name, usageLine, long string) *flag.FlagSet {
	fs := flag.NewFlagSet(name, flag.ExitOnError)
	fs.Usage = func() {
		w := fs.Output()
		fmt.Fprintf(w, "Usage: %s\n\n%s\n", usageLine, strings.TrimSpace(long))
		hasFlags := false
		fs.VisitAll(func(*flag.Flag) { hasFlags = true })
		if hasFlags {
			fmt.Fprintln(w, "\nFlags:")
			fs.PrintDefaults()
		}
	}
	return fs
}

// readSource returns the contents of path, or stdin when path is "-".
func readSource(path string) (string, error) {
	if path == "-" {
		b, err := io.ReadAll(os.Stdin)
		return string(b), err
	}
	b, err := os.ReadFile(path)
	return string(b), err
}

func cmdTokens(args []string) int {
	fs := newCommand("tokens", "gasm tokens <file>", `
Print the lexical token stream of FILE: position, token kind and text, one
token per line.  FILE may be "-" to read standard input.
`)
	fs.Parse(args)
	if fs.NArg() != 1 {
		fmt.Fprintln(os.Stderr, "usage: gasm tokens <file>")
		return 2
	}
	src, err := readSource(fs.Arg(0))
	if err != nil {
		fmt.Fprintln(os.Stderr, "gasm:", err)
		return 1
	}
	for _, tok := range lexer.Tokenize(src) {
		fmt.Printf("%s\t%s\t%q\n", tok.Pos, tok.Kind, tok.Text)
	}
	return 0
}

func cmdParse(args []string) int {
	fs := newCommand("parse", "gasm parse <file>", `
Parse FILE and report syntax errors on stderr.  On success, print how many
declarations and TEXT functions the file contains.  FILE may be "-" to read
standard input.
`)
	fs.Parse(args)
	if fs.NArg() != 1 {
		fmt.Fprintln(os.Stderr, "usage: gasm parse <file>")
		return 2
	}
	path := fs.Arg(0)
	src, err := readSource(path)
	if err != nil {
		fmt.Fprintln(os.Stderr, "gasm:", err)
		return 1
	}
	file, errs := parser.Parse(path, src)
	for _, e := range errs {
		fmt.Fprintf(os.Stderr, "%s: %v\n", path, e)
	}
	if len(errs) > 0 {
		return 1
	}
	funcs := 0
	for _, d := range file.Decls {
		if _, ok := d.(*ast.Text); ok {
			funcs++
		}
	}
	fmt.Printf("%s: OK — %d declarations, %d functions\n", path, len(file.Decls), funcs)
	return 0
}

func cmdFmt(args []string) int {
	fs := newCommand("fmt", "gasm fmt [-w] [path...]", `
Canonicalise the formatting of Plan 9 assembly sources: indentation, operand
spacing, per-function mnemonic alignment and blank-line layout (exactly one
blank line before each label, TEXT and GLOBL block).  Formatting is
idempotent and preserves every line, comments included.

With no paths — or a directory path — every .s file below it is reformatted
in place and the changed files are listed, the way go fmt does; "." and "_"
directories are skipped.  Explicit file paths print to stdout unless -w is
given.
`)
	write := fs.Bool("w", false, "write result to the source file")
	fs.Parse(args)
	// Like go fmt: with no arguments, or with a directory argument, every .s
	// file below the directory is formatted in place and the names of the
	// changed files are listed; explicit file arguments keep the -w / stdout
	// behaviour.
	paths := fs.Args()
	dirMode := len(paths) == 0
	if dirMode {
		paths = []string{"."}
	}
	var files []string
	for _, p := range paths {
		info, err := os.Stat(p)
		if err != nil {
			fmt.Fprintln(os.Stderr, "gasm:", err)
			return 1
		}
		if info.IsDir() {
			dirMode = true
			found, err := asmFiles(p)
			if err != nil {
				fmt.Fprintln(os.Stderr, "gasm:", err)
				return 1
			}
			files = append(files, found...)
			continue
		}
		files = append(files, p)
	}
	rc := 0
	for _, path := range files {
		src, err := readSource(path)
		if err != nil {
			fmt.Fprintln(os.Stderr, "gasm:", err)
			rc = 1
			continue
		}
		out := format.Source(path, src)
		if dirMode || *write {
			if out != src {
				if err := os.WriteFile(path, []byte(out), 0o644); err != nil {
					fmt.Fprintln(os.Stderr, "gasm:", err)
					rc = 1
					continue
				}
				if dirMode {
					fmt.Println(path)
				}
			}
			continue
		}
		fmt.Print(out)
	}
	return rc
}

// asmFiles collects the .s files below dir, skipping directories whose name
// starts with "." or "_" — as the go tooling does, which keeps .git and
// scratch or reference trees (e.g. _refs) untouched.
func asmFiles(dir string) ([]string, error) {
	var out []string
	err := filepath.WalkDir(dir, func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if d.IsDir() {
			if path != dir && (strings.HasPrefix(d.Name(), ".") || strings.HasPrefix(d.Name(), "_")) {
				return filepath.SkipDir
			}
			return nil
		}
		if strings.HasSuffix(d.Name(), ".s") {
			out = append(out, path)
		}
		return nil
	})
	return out, err
}

func cmdLint(args []string) int {
	fs := newCommand("lint", "gasm lint <file...>", `
Run the static checks over the given files and print diagnostics as
"file:line:col: severity: message [code]".  The exit status is non-zero when
an error-severity diagnostic is found; warnings (e.g. the register-clobber
audit) do not affect it.

Rules include unknown-instruction, operand-count, undefined-label,
duplicate-label, missing-ret, missing-textflag-include, abi-argsize,
unreachable-code, register-clobber and funcdata-pcdata.
`)
	disable := fs.String("disable", "", "comma-separated rule codes to disable")
	fs.Parse(args)
	if fs.NArg() == 0 {
		fmt.Fprintln(os.Stderr, "usage: gasm lint <file...>")
		return 2
	}
	disabled := map[string]bool{}
	for _, code := range strings.Split(*disable, ",") {
		if code = strings.TrimSpace(code); code != "" {
			disabled[code] = true
		}
	}
	hadError := false
	for _, path := range fs.Args() {
		src, err := readSource(path)
		if err != nil {
			fmt.Fprintln(os.Stderr, "gasm:", err)
			hadError = true
			continue
		}
		file, errs := parser.Parse(path, src)
		for _, e := range errs {
			fmt.Fprintf(os.Stderr, "%s: %v\n", path, e)
			hadError = true
		}
		diags := lint.File(file, lint.Config{Arch: arch.FromFilename(path), Disable: disabled})
		for _, d := range diags {
			fmt.Printf("%s:%d:%d: %s: %s [%s]\n", path, d.Pos.Line, d.Pos.Column, d.Severity, d.Message, d.Code)
			if d.Severity == lint.Error {
				hadError = true
			}
		}
	}
	if hadError {
		return 1
	}
	return 0
}

func cmdLSP(args []string) int {
	fs := newCommand("lsp", "gasm lsp", `
Run the language server over standard input/output: JSON-RPC 2.0 with
Content-Length framing.  Point an LSP-capable editor at the binary and
associate it with .s files; the target architecture is inferred from the file
suffix (_amd64.s, _arm64.s, _riscv64.s, _loong64.s).  Provides completion,
hover, document symbols, diagnostics and semantic-token highlighting.
`)
	fs.Parse(args)
	srv := lsp.New(os.Stdin, os.Stdout)
	if err := srv.Run(); err != nil {
		fmt.Fprintln(os.Stderr, "gasm lsp:", err)
		return 1
	}
	return 0
}

func cmdAsm(args []string) int {
	fs := newCommand("asm", "gasm asm [--format raw|elf|macho|goobj] [-p pkg] [-o out] <file>", `
Assemble FILE (amd64) without the Go toolchain: every TEXT function is
encoded to machine code — scalar, VEX/AVX2 and EVEX/AVX-512 instructions,
FP/SP frame mapping, local labels and file-local static symbols (GLOBL/DATA)
resolved RIP-relative — and printed as a hex dump.

With -o the output is written to a file instead.  The --format flag selects
what is written: raw (the default) concatenates the functions and the data
section into one self-consistent image; elf and macho emit a relocatable
object (.text/.data sections, a symbol table and one PC32 relocation per
static-symbol reference) that links with the system toolchain; goobj emits
the Go toolchain's own object format, which cmd/link consumes directly (it
requires -p, the package path, and the installed Go toolchain).
`)
	out := fs.String("o", "", "write the output to this file")
	format := fs.String("format", "raw", "output format: raw (concatenated image), elf, macho or goobj (Go object)")
	pkg := fs.String("p", "", "package path for --format goobj (qualifies the exported symbols)")
	fs.Parse(args)
	if fs.NArg() != 1 {
		fmt.Fprintln(os.Stderr, "usage: gasm asm [--format raw|elf|macho|goobj] [-p pkg] [-o out] <file>")
		return 2
	}
	path := fs.Arg(0)
	if arch.FromFilename(path) != arch.AMD64 {
		fmt.Fprintln(os.Stderr, "gasm asm: only amd64 is supported in this Phase 2 increment")
		return 1
	}
	src, err := readSource(path)
	if err != nil {
		fmt.Fprintln(os.Stderr, "gasm:", err)
		return 1
	}
	f, errs := parser.Parse(path, src)
	for _, e := range errs {
		fmt.Fprintf(os.Stderr, "%s: %v\n", path, e)
	}
	if len(errs) > 0 {
		return 1
	}

	img, err := asm.AssembleFile(f)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s: %v\n", path, err)
		return 1
	}
	if len(img.Funcs) == 0 {
		fmt.Fprintln(os.Stderr, "gasm asm: no assemblable TEXT functions found")
		return 1
	}
	for _, fn := range img.Funcs {
		code := img.Code[fn.Offset : fn.Offset+fn.Size]
		fmt.Printf("%s: %d bytes\n", fn.Name, fn.Size)
		for i := 0; i < len(code); i += 16 {
			end := i + 16
			if end > len(code) {
				end = len(code)
			}
			fmt.Printf("  %04x:", i)
			for _, b := range code[i:end] {
				fmt.Printf(" %02x", b)
			}
			fmt.Println()
		}
	}
	if len(img.Data) > 0 {
		fmt.Printf("data: %d bytes at 0x%x\n", len(img.Data), len(img.Code))
		for _, d := range f.Decls {
			g, ok := d.(*ast.Globl)
			if !ok || g.Name == nil || g.Name.Pseudo != "SB" {
				continue
			}
			size := 0
			if g.Size != nil && g.Size.Imm.HasVal {
				size = int(g.Size.Imm.Val)
			}
			fmt.Printf("  %s: %d bytes at 0x%x\n", g.Name.Name, size, img.Symbols[g.Name.Name])
		}
		for i := 0; i < len(img.Data); i += 16 {
			end := i + 16
			if end > len(img.Data) {
				end = len(img.Data)
			}
			fmt.Printf("  %04x:", len(img.Code)+i)
			for _, b := range img.Data[i:end] {
				fmt.Printf(" %02x", b)
			}
			fmt.Println()
		}
	}
	if *out != "" {
		var obj []byte
		var err error
		var kind string
		switch *format {
		case "raw":
			if len(img.Externals) > 0 {
				fmt.Fprintf(os.Stderr, "gasm asm: external symbol %q needs an object file (use --format elf or --format macho)\n", img.Externals[0])
				return 1
			}
			obj, kind = img.Bytes(), "raw image"
		case "elf":
			obj, err = img.ELFObject()
			kind = "ELF object"
		case "macho":
			obj, err = img.MachOObject()
			kind = "Mach-O object"
		case "goobj":
			obj, err = img.GOObject(*pkg, path)
			kind = "Go object"
		default:
			fmt.Fprintf(os.Stderr, "gasm asm: unknown format %q (want raw, elf, macho or goobj)\n", *format)
			return 2
		}
		if err != nil {
			fmt.Fprintln(os.Stderr, "gasm asm:", err)
			return 1
		}
		if err := os.WriteFile(*out, obj, 0o644); err != nil {
			fmt.Fprintln(os.Stderr, "gasm asm:", err)
			return 1
		}
		fmt.Printf("wrote %d bytes to %s (%s)\n", len(obj), *out, kind)
	}
	return 0
}