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

package asm

import (
	"fmt"
	"sort"
	"strconv"

	"sourcedock.dev/petrbalvin/gasm-devkit/ast"
)

// Image is an assembled file: the function bodies laid out in source order,
// followed by the file's static data section (GLOBL/DATA).  References to
// file-local static symbols are encoded RIP-relative and resolved within the
// image, so the raw bytes are self-consistent and executable at any base
// address; references to external symbols are recorded as relocations
// (Funcs[i].Relocs, Externals) and left unresolved, the object-file
// emitters turn them into linker relocations.
type Image struct {
	Code      []byte         // concatenated function bodies
	Data      []byte         // static data section
	Funcs     []FuncLayout   // function positions, in source order
	Symbols   map[string]int // static symbol → byte offset within the image
	DataSyms  []DataSymbol   // GLOBL symbols, in layout order
	Externals []string       // referenced but undefined symbols, sorted
}

// FuncLayout describes one assembled function within an Image.
type FuncLayout struct {
	Name    string
	Pkg     string // explicit package prefix ("" = the current package)
	Static  bool   // the <> marker: file-local, not exported
	Offset  int    // start offset within the image (== offset within Code)
	Size    int
	Args    int            // declared argument/result area (the TEXT size suffix)
	Frame   int            // local frame size (the TEXT $framesize)
	NoSplit bool           // the NOSPLIT flag
	SPWrite bool           // the SPWRITE flag: writes an arbitrary value to SP
	Line    int            // source line of the TEXT directive
	Labels  map[string]int // local labels, function-relative
	Relocs  []Reloc        // static-symbol references, in emission order
	Spadj   []SpadjStep    // stack-adjustment boundaries, ascending by PC
	Lines   []LineEntry    // source-line table: byte offset → source line
}

// SpadjStep is one stack-adjustment boundary: Value is the SP delta from the
// entry state in effect from PC (function-relative) until the next step.
type SpadjStep struct {
	PC    int
	Value int
}

// LineEntry maps a byte offset (function-relative) to a source line number.
type LineEntry struct {
	Offset int
	Line   int
}

// LineAt returns the source line number for the given function-relative byte
// offset, using a binary search on the line table.  Returns 0 if the offset
// is before the first instruction or the table is empty.
func (fl *FuncLayout) LineAt(offset int) int {
	if len(fl.Lines) == 0 {
		return 0
	}
	// Binary search: find the last entry with Offset <= offset.
	lo, hi := 0, len(fl.Lines)-1
	for lo < hi {
		mid := (lo + hi + 1) / 2
		if fl.Lines[mid].Offset <= offset {
			lo = mid
		} else {
			hi = mid - 1
		}
	}
	if fl.Lines[lo].Offset <= offset {
		return fl.Lines[lo].Line
	}
	return 0
}

// RelocKind Reloc is one static-symbol reference within a function body: the disp32
// field at Off (function-relative) must reach the symbol plus Addend,
// measured from After, the address just past the instruction.  An External
// relocation names a symbol no GLOBL in the file defines; the object-file
// emitters carry it into the output's relocation table.
// RelocKind discriminates the type of relocation needed.
type RelocKind int

const (
	RelPCRel32         RelocKind = iota // 32-bit PC-relative (amd64)
	RelCall                             // R_CALL: CALL to a function symbol (amd64)
	RelTLSLE                            // R_TLS_LE: local-exec TLS load, no symbol (amd64 guard)
	RelRISCVPCRELIType                  // R_RISCV_PCREL_ITYPE (AUIPC + I-type pair)
	RelRISCVPCRELSType                  // R_RISCV_PCREL_STYPE (AUIPC + S-type pair)
	RelRISCVJal                         // R_RISCV_JAL (J-type call)
	RelPCRelAbs                         // 32-bit absolute (R_RISCV_32)
	RelLoong64AddrHi                    // R_LOONG64_ADDR_HI (pcalau12i)
	RelLoong64AddrLo                    // R_LOONG64_ADDR_LO (addi.d/ld/st)
	RelArm64Addr                        // R_ADDRARM64 (ADRP + ADD pair)
	RelArm64Branch                      // R_CALLARM64 (BL instruction)
	RelArm64LDST64                      // R_ARM64_PCREL_LDST64 (ADRP + 64-bit LDR/STR pair)
	RelLoong64Branch                    // R_CALLLOONG64 (BL instruction)
)

type Reloc struct {
	Off      int
	After    int
	Name     string
	Addend   int64
	External bool
	Kind     RelocKind
}

// DataSymbol describes one GLOBL symbol laid out in the data section.
type DataSymbol struct {
	Name   string
	Pkg    string // explicit package prefix ("" = the current package)
	Offset int    // byte offset within Data
	Size   int
	Static bool // the <> marker: file-local, not exported
	Rodata bool // the RODATA flag: read-only data
	Dupok  bool // the DUPOK flag: duplicate-OK
}

// Bytes returns the whole image: code, then data.
func (img *Image) Bytes() []byte {
	out := make([]byte, 0, len(img.Code)+len(img.Data))
	out = append(out, img.Code...)
	return append(out, img.Data...)
}

// AssembleFile assembles every TEXT function of a parsed file and lays out
// its static symbols (GLOBL/DATA) in a data section behind the code.  Each
// reference to a file-local static symbol becomes a RIP-relative load whose
// displacement is resolved against that layout; a reference to a symbol no
// GLOBL defines is recorded as an external relocation (Externals) with its
// displacement left zero, the object-file emitters resolve it at link
// time, while the raw image (Bytes) cannot represent it.
func AssembleFile(f *ast.File) (*Image, error) {
	dataSyms, err := collectData(f)
	if err != nil {
		return nil, err
	}
	known := make(map[string]bool, len(dataSyms))
	for _, d := range dataSyms {
		known[d.name] = true
	}
	link := &linkInfo{symbols: known, allowExternal: true}

	img := &Image{Symbols: map[string]int{}}
	textOff := map[string]int{}
	type asmFunc struct {
		name    string
		patches []sbPatch
	}
	var funcs []asmFunc
	for _, d := range f.Decls {
		t, ok := d.(*ast.Text)
		if !ok {
			continue
		}
		code, patches, labels, steps, lines, err := assemble(t, link)
		if err != nil {
			return nil, fmt.Errorf("%s: %w", t.Name.Name, err)
		}
		fl := FuncLayout{
			Name:   t.Name.Name,
			Pkg:    t.Name.Pkg,
			Static: t.Name.Static,
			Offset: len(img.Code),
			Size:   len(code),
			Frame:  frameSize(t),
			Args:   argsSize(t),
			Line:   t.Pos().Line,
			Labels: labels,
			Lines:  lines,
		}
		for _, f := range t.Flags {
			switch f {
			case "NOSPLIT":
				fl.NoSplit = true
			case "SPWRITE":
				fl.SPWrite = true
			}
		}
		for _, s := range steps {
			fl.Spadj = append(fl.Spadj, SpadjStep{PC: s.pc, Value: s.value})
		}
		textOff[t.Name.Name] = len(img.Code)
		img.Funcs = append(img.Funcs, fl)
		img.Code = append(img.Code, code...)
		funcs = append(funcs, asmFunc{name: t.Name.Name, patches: patches})
	}

	// Lay out the data section behind the code, each symbol 16-aligned.
	dataStart := len(img.Code)
	for _, d := range dataSyms {
		if pos := dataStart + len(img.Data); pos != align16(pos) {
			img.Data = append(img.Data, make([]byte, align16(pos)-pos)...)
		}
		img.Symbols[d.name] = dataStart + len(img.Data)
		img.DataSyms = append(img.DataSyms, DataSymbol{
			Name:   d.name,
			Pkg:    d.pkg,
			Offset: len(img.Data),
			Size:   len(d.buf),
			Static: d.static,
			Rodata: d.rodata,
			Dupok:  d.dupok,
		})
		img.Data = append(img.Data, d.buf...)
	}

	// Resolve the RIP-relative displacements of file-local references now
	// that every address is known, and record every reference (resolved or
	// external) for the object-file emitters.
	externals := map[string]bool{}
	for i, fn := range funcs {
		base := img.Funcs[i].Offset
		code := img.Code[base : base+img.Funcs[i].Size]
		for _, p := range fn.patches {
			reloc := Reloc{Off: p.off, After: p.after, Name: p.name, Addend: p.addend, Kind: p.kind}
			if p.kind == RelTLSLE {
				// The TLS slot has no symbol: the linker fills the offset
				// from the runtime's TLS layout.
				img.Funcs[i].Relocs = append(img.Funcs[i].Relocs, reloc)
				continue
			}
			if imgOff, ok := img.Symbols[p.name]; ok {
				rel := int64(imgOff) + p.addend - int64(base+p.after)
				if rel < -1<<31 || rel >= 1<<31 {
					return nil, fmt.Errorf("%s: displacement to %q out of rel32 range", fn.name, p.name)
				}
				copy(code[p.off:p.off+4], le32(rel))
			} else if imgOff, ok := textOff[p.name]; ok {
				// A CALL to a TEXT function of the same file: resolve the
				// displacement against the function's layout position.
				rel := int64(imgOff) + p.addend - int64(base+p.after)
				if rel < -1<<31 || rel >= 1<<31 {
					return nil, fmt.Errorf("%s: displacement to %q out of rel32 range", fn.name, p.name)
				}
				copy(code[p.off:p.off+4], le32(rel))
			} else {
				reloc.External = true
				externals[p.name] = true
			}
			img.Funcs[i].Relocs = append(img.Funcs[i].Relocs, reloc)
		}
	}
	for name := range externals {
		img.Externals = append(img.Externals, name)
	}
	sort.Strings(img.Externals)
	return img, nil
}

// AssembleFileRISCV assembles every TEXT function of a parsed RISC-V file
// and lays out its static symbols (GLOBL/DATA) in a data section behind the
// code.  SB references in the code are encoded as AUIPC pairs with zero
// immediates; the object-file emitters record relocations for the linker.
func AssembleFileRISCV(f *ast.File) (*Image, error) {
	dataSyms, err := collectData(f)
	if err != nil {
		return nil, err
	}

	img := &Image{Symbols: map[string]int{}}
	for _, d := range f.Decls {
		t, ok := d.(*ast.Text)
		if !ok {
			continue
		}
		code, labels, relocs, lines, spadj, err := assembleRISCV(t)
		if err != nil {
			return nil, fmt.Errorf("%s: %w", t.Name.Name, err)
		}
		fl := FuncLayout{
			Name:   t.Name.Name,
			Pkg:    t.Name.Pkg,
			Static: t.Name.Static,
			Offset: len(img.Code),
			Size:   len(code),
			Frame:  frameSize(t),
			Args:   argsSize(t),
			Line:   t.Pos().Line,
			Labels: labels,
			Lines:  lines,
			Spadj:  spadj,
			Relocs: relocs,
		}
		for _, f := range t.Flags {
			switch f {
			case "NOSPLIT":
				fl.NoSplit = true
			case "SPWRITE":
				fl.SPWrite = true
			}
		}
		img.Funcs = append(img.Funcs, fl)
		img.Code = append(img.Code, code...)
	}

	// Lay out the data section behind the code, 16-aligned.
	dataStart := len(img.Code)
	for _, d := range dataSyms {
		pos := dataStart + len(img.Data)
		for pos%16 != 0 {
			img.Data = append(img.Data, 0)
			pos++
		}
		img.Symbols[d.name] = pos
		img.Data = append(img.Data, d.buf...)
		img.DataSyms = append(img.DataSyms, DataSymbol{
			Name:   d.name,
			Pkg:    d.pkg,
			Offset: len(img.Data) - len(d.buf), // relative to the data section
			Size:   d.size,
			Static: d.static,
			Rodata: d.rodata,
			Dupok:  d.dupok,
		})
	}

	markExternals(img, dataSyms)
	return img, nil
}

// AssembleFileLOONG64 assembles every TEXT function of a parsed loong64 file
// and lays out its static symbols (GLOBL/DATA) in a data section behind the
// code.  SB references in the code are encoded as pcalau12i pairs with zero
// immediates; the object-file emitters record R_LOONG64_ADDR_HI/LO
// relocations for the linker.
func AssembleFileLOONG64(f *ast.File) (*Image, error) {
	dataSyms, err := collectData(f)
	if err != nil {
		return nil, err
	}

	img := &Image{Symbols: map[string]int{}}
	for _, d := range f.Decls {
		t, ok := d.(*ast.Text)
		if !ok {
			continue
		}
		code, labels, relocs, lines, spadj, err := assembleLOONG64(t)
		if err != nil {
			return nil, fmt.Errorf("%s: %w", t.Name.Name, err)
		}
		fl := FuncLayout{
			Name:   t.Name.Name,
			Pkg:    t.Name.Pkg,
			Static: t.Name.Static,
			Offset: len(img.Code),
			Size:   len(code),
			Frame:  frameSize(t),
			Args:   argsSize(t),
			Line:   t.Pos().Line,
			Labels: labels,
			Lines:  lines,
			Spadj:  spadj,
			Relocs: relocs,
		}
		for _, f := range t.Flags {
			switch f {
			case "NOSPLIT":
				fl.NoSplit = true
			case "SPWRITE":
				fl.SPWrite = true
			}
		}
		img.Funcs = append(img.Funcs, fl)
		img.Code = append(img.Code, code...)
	}

	// Lay out the data section behind the code, 16-aligned.
	dataStart := len(img.Code)
	for _, d := range dataSyms {
		pos := dataStart + len(img.Data)
		for pos%16 != 0 {
			img.Data = append(img.Data, 0)
			pos++
		}
		img.Symbols[d.name] = pos
		img.Data = append(img.Data, d.buf...)
		img.DataSyms = append(img.DataSyms, DataSymbol{
			Name:   d.name,
			Pkg:    d.pkg,
			Offset: len(img.Data) - len(d.buf), // relative to the data section
			Size:   d.size,
			Static: d.static,
			Rodata: d.rodata,
			Dupok:  d.dupok,
		})
	}

	markExternals(img, dataSyms)
	return img, nil
}

// markExternals identifies relocations that reference symbols not defined in
// the file (neither a GLOBL/DATA symbol nor a TEXT function) and records them
// as external.  The non-amd64 architectures emit relocations for every SB
// reference; this post-processing step distinguishes file-local from external.
func markExternals(img *Image, dataSyms []dataSym) {
	known := make(map[string]bool, len(dataSyms)+len(img.Funcs))
	for _, d := range dataSyms {
		known[d.name] = true
	}
	for _, fn := range img.Funcs {
		known[fn.Name] = true
	}
	externals := map[string]bool{}
	for i := range img.Funcs {
		for j := range img.Funcs[i].Relocs {
			r := &img.Funcs[i].Relocs[j]
			if !known[r.Name] {
				r.External = true
				externals[r.Name] = true
			}
		}
	}
	for name := range externals {
		img.Externals = append(img.Externals, name)
	}
	sort.Strings(img.Externals)
}

// dataSym is one GLOBL symbol and its DATA initialiser.
type dataSym struct {
	name   string
	pkg    string
	buf    []byte
	size   int
	static bool
	rodata bool
	dupok  bool
}

// collectData gathers the file's static symbols (GLOBL) and their initial
// contents (DATA) into byte buffers, in declaration order.
func collectData(f *ast.File) ([]dataSym, error) {
	index := map[string]int{}
	var syms []dataSym
	for _, d := range f.Decls {
		switch dd := d.(type) {
		case *ast.Globl:
			if dd.Name == nil || dd.Name.Pseudo != "SB" {
				continue
			}
			name := dd.Name.Name
			if _, dup := index[name]; dup {
				return nil, fmt.Errorf("duplicate GLOBL %q", name)
			}
			size := 0
			if dd.Size != nil && dd.Size.Imm.HasVal {
				size = int(dd.Size.Imm.Val)
			}
			index[name] = len(syms)
			ds := dataSym{
				name:   name,
				pkg:    dd.Name.Pkg,
				buf:    make([]byte, size),
				size:   size,
				static: dd.Name.Static,
			}
			for _, f := range dd.Flags {
				switch f {
				case "RODATA":
					ds.rodata = true
				case "DUPOK":
					ds.dupok = true
				default:
					// Legacy numeric flag constants (runtime/textflag.h):
					// DUPOK is 2, RODATA is 8; combinations arrive as one
					// number (e.g. 10 = RODATA|DUPOK).
					if n, err := strconv.Atoi(f); err == nil {
						if n&2 != 0 {
							ds.dupok = true
						}
						if n&8 != 0 {
							ds.rodata = true
						}
					}
				}
			}
			syms = append(syms, ds)

		case *ast.Data:
			if dd.Name == nil || dd.Name.Pseudo != "SB" {
				continue
			}
			i, ok := index[dd.Name.Name]
			if !ok {
				return nil, fmt.Errorf("DATA %q: no matching GLOBL", dd.Name.Name)
			}
			if dd.Value == nil || !dd.Value.Imm.HasVal {
				return nil, fmt.Errorf("DATA %q: value must be an integer immediate", dd.Name.Name)
			}
			w := dd.Width
			switch w {
			case 1, 2, 4, 8:
			default:
				return nil, fmt.Errorf("DATA %q: invalid width %d (want 1, 2, 4 or 8)", dd.Name.Name, w)
			}
			off := dd.Name.Offset
			buf := syms[i].buf
			if off < 0 || off+int64(w) > int64(len(buf)) {
				return nil, fmt.Errorf("DATA %q+%d/%d exceeds GLOBL size %d", dd.Name.Name, off, w, len(buf))
			}
			v := dd.Value.Imm.Val
			if dd.Value.Imm.Neg {
				v = -v
			}
			for j := range w {
				buf[off+int64(j)] = byte(v >> (8 * j))
			}
		}
	}
	return syms, nil
}

// align16 rounds n up to the next multiple of 16.
func align16(n int) int {
	return (n + 15) &^ 15
}

// frameSize returns the local frame size declared on the TEXT directive.
func frameSize(t *ast.Text) int {
	if t.Frame != nil && t.Frame.Imm.HasVal {
		return int(t.Frame.Imm.Val)
	}
	return 0
}

// argsSize returns the argument/result area declared on the TEXT directive.
func argsSize(t *ast.Text) int {
	if t.Args != nil && t.Args.Imm.HasVal {
		return int(t.Args.Imm.Val)
	}
	return 0
}