package assembler

import (
	"fmt"
	"git.gammaspectra.live/WeebDataHoarder/compute-go/assembler/asm/arch"
	"git.gammaspectra.live/WeebDataHoarder/compute-go/assembler/obj"
	"git.gammaspectra.live/WeebDataHoarder/compute-go/assembler/obj/x86"
	"git.gammaspectra.live/WeebDataHoarder/compute-go/assembler/objabi"
	"git.gammaspectra.live/WeebDataHoarder/compute-go/assembler/src"
	memory "git.gammaspectra.live/WeebDataHoarder/compute-go/malloc"
	"git.gammaspectra.live/WeebDataHoarder/compute-go/types"
	"path"
	"reflect"
	"runtime"
	"slices"
)

const GlobalFunctionPrefix = "·"

// Builder allows you to assemble a series of instructions.
type Builder struct {
	pageAllocator memory.PageAllocator
	abi0          *ABIConfig
	abiInternal   *ABIConfig
	abiDef        *types.ABIDefinition

	ctxt *obj.Link
	arch *arch.Arch

	first *obj.Prog
	last  *obj.Prog

	// bulk allocator.
	block  []*obj.Prog
	used   int
	errors []error

	pendingLabels []string
	labels        map[string]*obj.Prog
}

// Root returns the first instruction.
func (b *Builder) Root() *obj.Prog {
	return b.first
}

func (b *Builder) Reset() {
	for _, prog := range b.block {
		*prog = obj.Prog{}
	}

	b.first = nil
	b.last = nil
	b.errors = nil
	b.pendingLabels = nil
	clear(b.labels)
	b.used = 0
}

// NewProg returns a new instruction structure for obj.ProgAlloc interface
func (b *Builder) NewProg() *obj.Prog {
	var p *obj.Prog

	if b.used >= len(b.block) {
		p = b.ctxt.NewProg()
		b.block = append(b.block, p)
		b.used++
	} else {
		p = b.block[b.used]
		b.used++
	}

	p.Ctxt = b.ctxt

	return p
}

func (b *Builder) GetLabelByProgram(p *obj.Prog) string {
	for k, l := range b.labels {
		if l == p {
			return k
		}
	}
	return ""
}

func (b *Builder) GetLabel(name string) *obj.Prog {
	return b.labels[name]
}

func (b *Builder) GetLabelLocation(name string) types.Address {
	return types.Address{
		Type: obj.TYPE_BRANCH,
		Val:  b.GetLabel(name),
	}
}

func (b *Builder) Label(name string) {
	if b.GetLabel(name) != nil || slices.Contains(b.pendingLabels, name) {
		panic("label exists")
	}
	b.pendingLabels = append(b.pendingLabels, name)
}

func (b *Builder) Last() *obj.Prog {
	return b.last
}

// AddInstruction adds an instruction to the list of instructions
// to be assembled.
func (b *Builder) AddInstruction(p *obj.Prog, skipn ...int) {

	//todo: bring skip from caller
	skip := 1
	if len(skipn) > 0 {
		skip += skipn[0]
	}
	_, file, line, ok := runtime.Caller(skip)
	if ok {
		p.Pos = b.XPos(file, line)
	}

	if b.first == nil {
		b.first = p
		b.last = p
	} else {
		b.last.Link = p
		b.last = p
	}

	for len(b.pendingLabels) > 0 {
		label := b.pendingLabels[len(b.pendingLabels)-1]
		b.labels[label] = p
		b.pendingLabels = b.pendingLabels[:len(b.pendingLabels)-1]
	}
}

func (b *Builder) XPos(file string, line int) src.XPos {
	fBase := src.NewFileBase(file, objabi.AbsFile(path.Dir(file), file, ""))
	return b.ctxt.PosTable.XPos(src.MakePos(fBase, uint(line), 0))
}

type ListingEntry struct {
	Prog     *obj.Prog
	Printing string
	Code     []byte
}

func FunctionArgs[F any](b *Builder, abi obj.ABI) *ABIParamResultInfo {
	var f F
	return b.FuncArgs(reflect.TypeOf(f), abi)
}

func (b *Builder) FuncParam(params *ABIParamResultInfo, n int, name string) []types.Addressable {
	param := params.InParam(n)
	if name == "" {
		name = param.Name
	}

	if len(param.Registers) > 0 {
		r := make([]types.Addressable, len(param.Registers))
		for i, reg := range param.Registers {
			if int(reg) >= len(b.abiDef.ABIIntegerRegisters) {
				r[i] = types.Register(b.abiDef.ABIFloatRegisters[int(reg)-len(b.abiDef.ABIIntegerRegisters)])
			} else {
				r[i] = types.Register(b.abiDef.ABIIntegerRegisters[int(reg)])
			}
		}
		return r
	} else {
		return []types.Addressable{
			types.Address{
				Type:   obj.TYPE_MEM,
				Reg:    x86.REG_NONE,
				Offset: param.FrameOffset(params),
				Name:   obj.NAME_PARAM,
				Sym: &obj.LSym{
					Name: name,
				},
			},
		}
	}
}

func (b *Builder) funcParam(params *ABIParamResultInfo, param *ABIParamAssignment, name string) []types.Addressable {
	if name == "" {
		name = param.Name
	}

	if len(param.Registers) > 0 {
		r := make([]types.Addressable, len(param.Registers))
		for i, reg := range param.Registers {
			if int(reg) >= len(b.abiDef.ABIIntegerRegisters) {
				r[i] = types.Register(b.abiDef.ABIFloatRegisters[int(reg)-len(b.abiDef.ABIIntegerRegisters)])
			} else {
				r[i] = types.Register(b.abiDef.ABIIntegerRegisters[int(reg)])
			}
		}
		return r
	} else {
		return []types.Addressable{
			types.Address{
				Type:   obj.TYPE_MEM,
				Reg:    x86.REG_NONE,
				Offset: param.FrameOffset(params),
				Name:   obj.NAME_PARAM,
				Sym: &obj.LSym{
					Name: name,
				},
			},
		}
	}
}

func (b *Builder) FuncArgs(fType reflect.Type, abi obj.ABI) *ABIParamResultInfo {
	if fType.Kind() != reflect.Func {
		panic("not a function")
	}

	switch abi {
	case obj.ABI0:
		return b.abi0.ABIAnalyze(NewIType(fType, "", b.arch))
	case obj.ABIInternal:
		return b.abiInternal.ABIAnalyze(NewIType(fType, "", b.arch))
	default:
		panic("unknown ABI")
	}
}

func (b *Builder) Function(name string, params *ABIParamResultInfo, frameSize, textFlag int, listing *[]ListingEntry) (asmCode []byte, errs []error) {

	nameAddr := obj.Addr{
		Type: obj.TYPE_MEM,
		Name: obj.NAME_EXTERN,
		Sym: &obj.LSym{
			Name: name,
		},
	}

	_, file, line, _ := runtime.Caller(1)
	xpos := b.XPos(file, line)

	b.ctxt.InitTextSym(nameAddr.Sym, textFlag, xpos)
	if params != nil {
		nameAddr.Sym.SetABI(params.config.Which())
	}

	prog := &obj.Prog{
		Ctxt: b.ctxt,
		As:   obj.ATEXT,
		Pos:  xpos,
		From: nameAddr,
		To: obj.Addr{
			Type:   obj.TYPE_TEXTSIZE,
			Offset: int64(frameSize),
			// Argsize set below.
		},
		Link: b.Root(),
	}
	nameAddr.Sym.Func().Text = prog
	if params != nil {
		prog.To.Val = int32(params.ArgWidth())
	} else {
		prog.To.Val = int32(0)
	}

	b.errors = nil

	plist := &obj.Plist{
		Firstpc: prog,
		Curfn:   nil,
	}

	obj.Flushplist(b.ctxt, plist, b.NewProg)
	if b.ctxt.Errors > 0 {
		return nil, b.errors
	}

	top := plist.Firstpc
	var text *obj.LSym

	//produce asm output
	if listing != nil {
		for p := top; p != nil; p = p.Link {
			if p.As == obj.ATEXT {
				text = p.From.Sym
			}

			if text == nil {
				b.errors = append(b.errors, fmt.Errorf("%s: instruction outside TEXT", p))
				continue
			}
			size := int64(len(text.P)) - p.Pc
			if p.Link != nil {
				size = p.Link.Pc - p.Pc
			} else if p.Isize != 0 {
				size = int64(p.Isize)
			}
			var code []byte
			if p.Pc < int64(len(text.P)) {
				code = text.P[p.Pc:]
				if size < int64(len(code)) {
					code = code[:size]
				}
			}

			*listing = append(*listing, ListingEntry{
				Prog:     p,
				Printing: p.InstructionString(),
				Code:     code,
			})
		}
	}

	return nameAddr.Sym.P, nil
}

// NewBuilder constructs an assembler for the given architecture.
func NewBuilder(abi *types.ABIDefinition) (*Builder, error) {
	//TODO: add custom instruction mappings!

	b := &Builder{
		arch:        abi.Arch,
		block:       make([]*obj.Prog, 0, 64),
		labels:      make(map[string]*obj.Prog),
		abi0:        NewABIConfig(0, 0, abi.Arch, obj.ABI0),
		abiInternal: NewABIConfig(len(abi.ABIIntegerRegisters), len(abi.ABIFloatRegisters), abi.Arch, obj.ABIInternal),
		abiDef:      abi,
	}

	ctxt := obj.Linknew(abi.Arch.LinkArch)
	ctxt.IsAsm = true
	err := ctxt.Headtype.Set(runtime.GOOS)
	if err != nil {
		return nil, err
	}
	//todo
	ctxt.Pkgpath = "internal"
	ctxt.DiagFunc = func(in string, args ...interface{}) {
		b.errors = append(b.errors, fmt.Errorf(in, args...))
	}
	abi.Arch.Init(ctxt)
	b.ctxt = ctxt

	return b, nil
}