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

package main

import (
	"fmt"
	"os"
	"sort"
	"strings"

	"sourcedock.dev/petrbalvin/gasm-devkit/arch"
	"sourcedock.dev/petrbalvin/gasm-devkit/disasm"
	"sourcedock.dev/petrbalvin/gasm-devkit/parser"
)

// cmdDis disassembles machine code: either a raw binary (standard input with
// "-") whose architecture is given with -a, or a .s file, which is assembled
// first so the listing shows the real function and label layout.
func cmdDis(args []string) int {
	fs := newCommand("dis", "gasm dis [-a arch] <file>", `
Disassemble machine code to instruction text (via golang.org/x/arch).

With a .s file, the file is assembled first and the listing follows the
real layout: one block per TEXT function, local labels printed at their
offsets.  The architecture comes from the file name suffix, or from -a.

With any other file, or "-" for standard input, the bytes are disassembled
linearly and -a selects the architecture (amd64, arm64, riscv64 or
loong64).
`)
	archName := fs.String("a", "", "architecture for raw input: amd64, arm64, riscv64 or loong64")
	fs.Parse(args)
	if fs.NArg() != 1 {
		fmt.Fprintln(os.Stderr, "usage: gasm dis [-a arch] <file>")
		return 2
	}
	path := fs.Arg(0)
	var target arch.Arch
	if *archName != "" {
		var err error
		target, err = auditArch(*archName)
		if err != nil {
			fmt.Fprintf(os.Stderr, "gasm dis: %v\n", err)
			return 2
		}
	}

	if strings.HasSuffix(path, ".s") {
		if target == arch.Unknown {
			target = arch.FromFilename(path)
		}
		if target == arch.Unknown {
			fmt.Fprintln(os.Stderr, "gasm dis: cannot infer the architecture from the file name; use -a")
			return 2
		}
		return disSource(path, target)
	}

	if target == arch.Unknown {
		fmt.Fprintln(os.Stderr, "gasm dis: raw input needs -a (amd64, arm64, riscv64 or loong64)")
		return 2
	}
	src, err := readSource(path)
	if err != nil {
		fmt.Fprintln(os.Stderr, "gasm dis:", err)
		return 1
	}
	printListing(target, []byte(src), 0, nil)
	return 0
}

// disSource assembles a .s file and prints one listing block per function.
func disSource(path string, target arch.Arch) int {
	src, err := readSource(path)
	if err != nil {
		fmt.Fprintln(os.Stderr, "gasm dis:", 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 := assembleFile(target, f)
	if err != nil {
		fmt.Fprintf(os.Stderr, "gasm dis: %v\n", err)
		return 1
	}
	if len(img.Funcs) == 0 {
		fmt.Fprintln(os.Stderr, "gasm dis: 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)
		labels := make(map[int][]string, len(fn.Labels))
		for name, off := range fn.Labels {
			labels[off] = append(labels[off], name)
		}
		for off := range labels {
			sort.Strings(labels[off])
		}
		printListing(target, code, uint64(fn.Offset), labels)
	}
	if len(img.Data) > 0 {
		fmt.Printf("data: %d bytes at 0x%x\n", len(img.Data), len(img.Code))
	}
	return 0
}

// printListing decodes code linearly from offset base, printing label lines
// (label name to offset within the block) as they are reached.
func printListing(a arch.Arch, code []byte, base uint64, labels map[int][]string) {
	pc := 0
	for pc < len(code) {
		for _, name := range labels[pc] {
			fmt.Printf("%s:\n", name)
		}
		ins, err := disasm.Decode(a, code[pc:], base+uint64(pc))
		if err != nil {
			break
		}
		end := min(pc+ins.Len, len(code))
		fmt.Printf("  %04x: %-16s %s\n", base+uint64(pc), hexBytes(code[pc:end]), ins.Text)
		if ins.Len <= 0 {
			break
		}
		pc += ins.Len
	}
}

// hexBytes renders up to 8 bytes as contiguous hex.
func hexBytes(b []byte) string {
	var sb strings.Builder
	for i, c := range b {
		if i == 8 {
			break
		}
		if i > 0 {
			sb.WriteByte(' ')
		}
		fmt.Fprintf(&sb, "%02x", c)
	}
	return sb.String()
}