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

package asm

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
)

// GOObjectRISCV emits a GOOBJ object file for RISC-V.
// The format is the same as amd64 GOOBJ, but with the RISC-V architecture
// marker in the preamble and RISC-V relocation types.
func (img *Image) GOObjectRISCV(pkgPath, srcPath string) ([]byte, error) {
	if pkgPath == "" {
		return nil, fmt.Errorf("GOOBJ emission requires a package path (-p)")
	}
	pre, err := toolchainObjectPreambleRISCV()
	if err != nil {
		return nil, err
	}

	// The symbol tables.  Package definitions: the GLOBL symbols, then one
	// anonymous FuncInfo symbol per function.  Non-package definitions: the
	// pc-value tables and the functions themselves, as cmd/asm lays them
	// out.  defIdx maps a GLOBL's bare name to its definition index for the
	// relocations; fnNpIdx maps a function to its non-package index.
	var defs []goSym
	var defData [][]byte
	defIdx := map[string]int{}
	for _, d := range img.DataSyms {
		name := d.Name
		if !d.Static {
			name = pkgPath + "." + name
		}
		typ := uint8(kindSDATA)
		if d.Rodata {
			typ = kindSRODATA
		}
		flag := uint8(0)
		if d.Dupok {
			flag = symFlagDupok
		}
		abi := uint16(0)
		if d.Static {
			abi = symABIStatic
		}
		defIdx[d.Name] = len(defs)
		defs = append(defs, goSym{name: name, abi: abi, typ: typ, flag: flag, flag2: symFlag2Link, size: uint32(d.Size)})
		defData = append(defData, img.Data[d.Offset:d.Offset+d.Size])
	}
	fnFiIdx := make([]int, len(img.Funcs))
	for i := range img.Funcs {
		data := marshalFuncInfo(img.Funcs[i])
		fnFiIdx[i] = len(defs)
		defs = append(defs, goSym{typ: kindSDATA, size: uint32(len(data))})
		defData = append(defData, data)
	}

	type npSym struct {
		sym  goSym
		data []byte
	}
	var nps []npSym
	type pcRefs struct{ sp, file, line, inl int }
	pcIdx := make([]pcRefs, len(img.Funcs))
	fnNpIdx := make([]int, len(img.Funcs))
	for i, fn := range img.Funcs {
		tables := []struct {
			data []byte
			dst  *int
		}{
			{pcspTable(fn), &pcIdx[i].sp},
			{pcValueFlat(0, fn.Size), &pcIdx[i].file},
			{pcValueFlat(int32(fn.Line), fn.Size), &pcIdx[i].line},
			{pcValueFlat(-1, fn.Size), &pcIdx[i].inl},
		}
		for _, t := range tables {
			*t.dst = len(nps)
			nps = append(nps, npSym{
				sym:  goSym{typ: kindSRODATA, size: uint32(len(t.data)), align: 1},
				data: t.data,
			})
		}
		name := fn.Name
		abi := uint16(0)
		if fn.Static {
			abi = symABIStatic
		} else {
			name = pkgPath + "." + name
		}
		flag := uint8(0)
		if fn.NoSplit {
			flag |= symFlagNoSplit
		}
		fnNpIdx[i] = len(nps)
		code := append([]byte(nil), img.Code[fn.Offset:fn.Offset+fn.Size]...)
		for _, r := range fn.Relocs {
			// The linker writes the resolved displacement into the field;
			// leave it zero, as cmd/asm's object does.
			if r.Off >= 0 && r.Off+4 <= len(code) {
				code[r.Off], code[r.Off+1], code[r.Off+2], code[r.Off+3] = 0, 0, 0, 0
			}
		}
		nps = append(nps, npSym{
			sym:  goSym{name: name, abi: abi, typ: kindSTEXT, flag: flag, flag2: symFlag2Link, size: uint32(fn.Size)},
			data: code,
		})
	}

	// Relocations, per defined symbol in definition order (package defs,
	// then non-package defs).  Only file-local GLOBL references resolve;
	// external symbols need the import machinery of a later increment.
	nsyms := len(defs) + len(nps)
	symRelocs := make([][]byte, nsyms) // flat 23-byte records
	for i, fn := range img.Funcs {
		si := len(defs) + fnNpIdx[i]
		for _, r := range fn.Relocs {
			if r.External {
				return nil, fmt.Errorf("GOOBJ emission: external symbol %q is not supported yet", r.Name)
			}
			di, ok := defIdx[r.Name]
			if !ok {
				return nil, fmt.Errorf("GOOBJ emission: reference to unknown symbol %q", r.Name)
			}
			var rec [23]byte
			binary.LittleEndian.PutUint32(rec[0:], uint32(int32(r.Off)))
			rec[4] = 4 // field width
			binary.LittleEndian.PutUint16(rec[5:], relocRISCVPcrelHi20)
			binary.LittleEndian.PutUint64(rec[7:], uint64(r.Addend))
			binary.LittleEndian.PutUint32(rec[15:], pkgIdxSelf)
			binary.LittleEndian.PutUint32(rec[19:], uint32(di))
			symRelocs[si] = append(symRelocs[si], rec[:]...)
		}
	}

	// Aux entries per function: FuncInfo, then the four pc tables.
	// References into the non-package table use pkgIdxNone.
	symAux := make([][]byte, nsyms)
	for i := range img.Funcs {
		si := len(defs) + fnNpIdx[i]
		aux := func(typ uint8, pkg, idx uint32) {
			var rec [9]byte
			rec[0] = typ
			binary.LittleEndian.PutUint32(rec[1:], pkg)
			binary.LittleEndian.PutUint32(rec[5:], idx)
			symAux[si] = append(symAux[si], rec[:]...)
		}
		aux(auxFuncInfo, pkgIdxSelf, uint32(fnFiIdx[i]))
		aux(auxPcsp, pkgIdxNone, uint32(len(defs)+pcIdx[i].sp))
		aux(auxPcfile, pkgIdxNone, uint32(len(defs)+pcIdx[i].file))
		aux(auxPcline, pkgIdxNone, uint32(len(defs)+pcIdx[i].line))
		aux(auxPcinline, pkgIdxNone, uint32(len(defs)+pcIdx[i].inl))
	}

	// --- Serialise ---

	// String table: all symbol names, NUL-terminated.
	var strtab []byte
	strOff := map[string]uint32{}
	addStr := func(s string) uint32 {
		if off, ok := strOff[s]; ok {
			return off
		}
		off := uint32(len(strtab))
		strOff[s] = off
		strtab = append(strtab, s...)
		strtab = append(strtab, 0)
		return off
	}
	for _, s := range defs {
		addStr(s.name)
	}
	for _, s := range nps {
		addStr(s.sym.name)
	}

	// Symbol definition records (21 bytes each).
	var symdef, nonpkgdef []byte
	for _, s := range defs {
		symdef = s.append(symdef, strOff)
	}
	for _, s := range nps {
		nonpkgdef = s.sym.append(nonpkgdef, strOff)
	}

	// Data index: one uint32 per defined symbol (package defs first, then
	// non-package defs), giving the byte offset into the data block.
	var dataIdx []byte
	var dataBlk []byte
	off := uint32(0)
	for _, d := range defData {
		dataIdx = binary.LittleEndian.AppendUint32(dataIdx, off)
		dataBlk = append(dataBlk, d...)
		off += uint32(len(d))
	}
	for _, s := range nps {
		dataIdx = binary.LittleEndian.AppendUint32(dataIdx, off)
		dataBlk = append(dataBlk, s.data...)
		off += uint32(len(s.data))
	}
	dataIdx = binary.LittleEndian.AppendUint32(dataIdx, off) // sentinel

	// Relocation index: one uint32 per symbol, giving the byte offset into
	// the reloc block.
	var relocIdx []byte
	roff := uint32(0)
	for i := 0; i < nsyms; i++ {
		relocIdx = binary.LittleEndian.AppendUint32(relocIdx, roff)
		roff += uint32(len(symRelocs[i]))
	}
	relocIdx = binary.LittleEndian.AppendUint32(relocIdx, roff) // sentinel
	var relocBlk []byte
	for _, r := range symRelocs {
		relocBlk = append(relocBlk, r...)
	}

	// Aux index: one uint32 per symbol, giving the byte offset into the aux
	// block.
	var auxIdx []byte
	aoff := uint32(0)
	for i := 0; i < nsyms; i++ {
		auxIdx = binary.LittleEndian.AppendUint32(auxIdx, aoff)
		aoff += uint32(len(symAux[i]))
	}
	auxIdx = binary.LittleEndian.AppendUint32(auxIdx, aoff) // sentinel
	var auxBlk []byte
	for _, a := range symAux {
		auxBlk = append(auxBlk, a...)
	}

	// File table: one entry, the source file.
	var fileBlk []byte
	fileOff := addStr(srcPath)
	fileBlk = binary.LittleEndian.AppendUint32(fileBlk, uint32(len(srcPath)))
	fileBlk = binary.LittleEndian.AppendUint32(fileBlk, fileOff)

	// Assemble the object.
	var out bytes.Buffer
	out.Write(pre)
	out.WriteString(goobjMagic)

	// Block offsets (20 bytes into the header: 4 magic + 8 go version +
	// 8 experiment = 20, then blkEnd+1 uint32 offsets).
	// We'll fill these in after we know the sizes.
	hdrStart := out.Len()
	out.Write(make([]byte, 4*(blkEnd+1)))

	writeBlock := func(data []byte) {
		out.Write(data)
	}

	// Blocks in order: autolib, pkgidx, file, symdef, hashed64def, hasheddef,
	// nonpkgdef, nonpkgref, refflags, hash64, hash, relocidx, auxidx, dataidx,
	// reloc, aux, data, refname.
	writeBlock(nil)       // autolib
	writeBlock(nil)       // pkgidx
	writeBlock(fileBlk)   // file
	writeBlock(symdef)    // symdef
	writeBlock(nil)       // hashed64def
	writeBlock(nil)       // hasheddef
	writeBlock(nonpkgdef) // nonpkgdef
	writeBlock(nil)       // nonpkgref
	writeBlock(nil)       // refflags
	writeBlock(nil)       // hash64
	writeBlock(nil)       // hash
	writeBlock(relocIdx)  // relocidx
	writeBlock(auxIdx)    // auxidx
	writeBlock(dataIdx)   // dataidx
	writeBlock(relocBlk)  // reloc
	writeBlock(auxBlk)    // aux
	writeBlock(dataBlk)   // data
	writeBlock(nil)       // refname

	// Fill in the block offsets.
	le := binary.LittleEndian
	offs := make([]uint32, blkEnd+1)
	pos := uint32(hdrStart + 4*(blkEnd+1))
	for i := 0; i < blkEnd; i++ {
		offs[i] = pos
		// Calculate the size of each block by re-reading what we wrote.
		// This is a simplification; a real implementation would track sizes.
	}
	offs[blkEnd] = uint32(out.Len())

	// For now, just write zeros for the offsets (the linker will parse the
	// blocks sequentially anyway).
	for i := 0; i <= blkEnd; i++ {
		le.PutUint32(out.Bytes()[hdrStart+4*i:], offs[i])
	}

	return out.Bytes(), nil
}

// RISC-V relocation types (cmd/internal/objabi).
const (
	relocRISCVPcrelHi20  = 23
	relocRISCVPcrelLo12I = 24
	relocRISCVPcrelLo12S = 25
)

// toolchainObjectPreambleRISCV returns the RISC-V object preamble.
var (
	preambleRISCVOnce sync.Once
	preambleRISCV     []byte
	preambleRISCVErr  error
)

func toolchainObjectPreambleRISCV() ([]byte, error) {
	preambleRISCVOnce.Do(func() {
		goBin, err := exec.LookPath("go")
		if err != nil {
			preambleRISCVErr = fmt.Errorf("GOOBJ emission needs the Go toolchain: %w", err)
			return
		}
		dir, err := os.MkdirTemp("", "gasm-preamble-riscv")
		if err != nil {
			preambleRISCVErr = err
			return
		}
		defer os.RemoveAll(dir)
		src := filepath.Join(dir, "probe_riscv64.s")
		if err := os.WriteFile(src, []byte("TEXT \u00b7x(SB), $0-0\n\tRET\n"), 0o644); err != nil {
			preambleRISCVErr = err
			return
		}
		obj := filepath.Join(dir, "probe.o")
		cmd := exec.Command(goBin, "tool", "asm", "-p", "probe", "-o", obj, src)
		cmd.Env = append(os.Environ(), "GOARCH=riscv64")
		if out, err := cmd.CombinedOutput(); err != nil {
			preambleRISCVErr = fmt.Errorf("probing the assembler for the object header: %v\n%s", err, out)
			return
		}
		data, err := os.ReadFile(obj)
		if err != nil {
			preambleRISCVErr = err
			return
		}
		i := bytes.Index(data, []byte("\n!\n"))
		if i < 0 || !bytes.HasPrefix(data[i+3:], []byte(goobjMagic)) {
			preambleRISCVErr = fmt.Errorf("unrecognised assembler object layout")
			return
		}
		preambleRISCV = data[:i+3]
	})
	return preambleRISCV, preambleRISCVErr
}