package codegen

import (
	"slices"

	"git.urbach.dev/cli/q/src/cpu"
	"git.urbach.dev/cli/q/src/ssa"
	"git.urbach.dev/cli/q/src/token"
)

// findFreeRegister finds a free register for the given value.
func (f *Function) findFreeRegister(step *Step) cpu.Register {
	usedRegisters := bitSet(0)

	if f.needsFramePointer {
		usedRegisters.Set(f.CPU.FramePointer)
	}

	switch value := step.Value.(type) {
	case *ssa.BinaryOp:
		for _, operand := range value.Inputs() {
			if !f.arch.operandConflict(value, operand) {
				continue
			}

			operandStep := f.ValueToStep[operand]

			if operandStep.Register != -1 {
				usedRegisters.Set(operandStep.Register)
			}
		}
	case *ssa.Phi:
		for index, pre := range step.Block.Predecessors {
			region := f.BlockToRegion[pre]
			last := f.Steps[region.End-1]
			incoming := value.Arguments[index]

			for _, live := range last.Live {
				if live.Register != -1 && live.Value != incoming {
					usedRegisters.Set(live.Register)
				}
			}
		}
	}

	for _, current := range f.Steps {
		// These checks need to happen regardless of whether the value is alive after execution.
		// If it is used as an operand, the operand restrictions of the architecture apply.
		binaryOp, isBinaryOp := current.Value.(*ssa.BinaryOp)

		if isBinaryOp && current.Register != -1 && f.arch.operandConflict(binaryOp, step.Value) {
			usedRegisters.Set(current.Register)
		}

		if isBinaryOp && !binaryOp.Op.IsComparison() {
			switch binaryOp.Op {
			case token.Div, token.Mod:
				if binaryOp.Right == step.Value {
					for _, reg := range f.CPU.DivisorRestricted {
						usedRegisters.Set(reg)
					}
				}
			case token.Shl, token.Shr:
				if current == step {
					for _, reg := range f.CPU.ShiftRestricted {
						usedRegisters.Set(reg)
					}
				}
			}
		}

		// If it's not alive in this step, ignore it.
		if !slices.Contains(current.Live, step) {
			continue
		}

		// Mark all the neighbor registers that are alive
		// at the same time as used.
		for _, live := range current.Live {
			if live.Register == -1 {
				continue
			}

			switch instr := live.Value.(type) {
			case *ssa.Field:
				_, isFieldFromCall := instr.Tuple.(*ssa.Call)

				if isFieldFromCall && live.Index > current.Index {
					usedRegisters.Set(f.CPU.Call.Out[instr.Index])
				}
			case *ssa.Parameter:
				if live.Index > current.Index {
					usedRegisters.Set(f.CPU.Call.In[instr.Index])
				}
			}

			usedRegisters.Set(live.Register)
		}

		// Ignore the definition itself.
		if current == step {
			continue
		}

		// Find all the registers that this instruction
		// would clobber and mark them as used.
		for _, reg := range f.clobberedRegisters(current) {
			usedRegisters.Set(reg)
		}
	}

	// Pick one of the register hints if possible.
	for _, reg := range step.Hints {
		if !usedRegisters.Has(reg) {
			return reg
		}
	}

	// Pick a general purpose register that's not used yet.
	for _, reg := range f.CPU.General {
		if !usedRegisters.Has(reg) {
			return reg
		}
	}

	// Pick a virtual register.
	for reg := f.CPU.MaxRegisters; reg <= 63; reg++ {
		if !usedRegisters.Has(reg) {
			stackOffset := uint(reg-f.CPU.MaxRegisters) * 8

			if stackOffset+8 > f.stackSize {
				f.stackSize = stackOffset + 8
			}

			return reg
		}
	}

	panic("no free registers")
}