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

package asm

import (
	"fmt"

	"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).  Static-symbol
// references are encoded RIP-relative and resolved within the image, so the
// bytes are self-consistent and executable at any base address.
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
}

// FuncLayout describes one assembled function within an Image.
type FuncLayout struct {
	Name   string
	Offset int // start offset within the image (== offset within Code)
	Size   int
	Labels map[string]int // local labels, function-relative
}

// 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
// static-symbol reference becomes a RIP-relative load whose displacement is
// resolved against that layout.  External (non-file-local) symbol references
// are rejected: they need object-file emission.
func AssembleFile(f *ast.File) (*Image, error) {
	syms, order, err := collectData(f)
	if err != nil {
		return nil, err
	}
	known := make(map[string]bool, len(syms))
	for name := range syms {
		known[name] = true
	}
	link := &linkInfo{symbols: known}

	img := &Image{Symbols: 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, err := assemble(t, link)
		if err != nil {
			return nil, fmt.Errorf("%s: %w", t.Name.Name, err)
		}
		img.Funcs = append(img.Funcs, FuncLayout{
			Name:   t.Name.Name,
			Offset: len(img.Code),
			Size:   len(code),
			Labels: labels,
		})
		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 _, name := range order {
		if pos := dataStart + len(img.Data); pos != align16(pos) {
			img.Data = append(img.Data, make([]byte, align16(pos)-pos)...)
		}
		img.Symbols[name] = dataStart + len(img.Data)
		img.Data = append(img.Data, syms[name]...)
	}

	// Resolve the RIP-relative displacements now that every address is known.
	for i, fn := range funcs {
		base := img.Funcs[i].Offset
		code := img.Code[base : base+img.Funcs[i].Size]
		for _, p := range fn.patches {
			rel := int64(img.Symbols[p.name]) + 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))
		}
	}
	return img, nil
}

// collectData gathers the file's static symbols (GLOBL) and their initial
// contents (DATA) into byte buffers, in declaration order.
func collectData(f *ast.File) (map[string][]byte, []string, error) {
	syms := map[string][]byte{}
	var order []string
	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 := syms[name]; dup {
				return nil, nil, fmt.Errorf("duplicate GLOBL %q", name)
			}
			size := 0
			if dd.Size != nil && dd.Size.Imm.HasVal {
				size = int(dd.Size.Imm.Val)
			}
			syms[name] = make([]byte, size)
			order = append(order, name)

		case *ast.Data:
			if dd.Name == nil || dd.Name.Pseudo != "SB" {
				continue
			}
			buf, ok := syms[dd.Name.Name]
			if !ok {
				return nil, nil, fmt.Errorf("DATA %q: no matching GLOBL", dd.Name.Name)
			}
			if dd.Value == nil || !dd.Value.Imm.HasVal {
				return nil, 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, nil, fmt.Errorf("DATA %q: invalid width %d (want 1, 2, 4 or 8)", dd.Name.Name, w)
			}
			off := dd.Name.Offset
			if off < 0 || off+int64(w) > int64(len(buf)) {
				return nil, 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 i := 0; i < w; i++ {
				buf[off+int64(i)] = byte(v >> (8 * i))
			}
		}
	}
	return syms, order, nil
}

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