package optimizer

import "strings"

type regInfo struct {
	val int    // -1 = unknown, else the immediate value known to be in reg
	src string // variable name if loaded from memory (or "" for immediates)
}

type regState struct {
	a regInfo
	x regInfo
	y regInfo
}

func newRegState() *regState {
	return &regState{
		a: regInfo{val: -1},
		x: regInfo{val: -1},
		y: regInfo{val: -1},
	}
}

func (rs *regState) reset() {
	rs.a = regInfo{val: -1}
	rs.x = regInfo{val: -1}
	rs.y = regInfo{val: -1}
}

func (rs *regState) resetA() { rs.a = regInfo{val: -1} }
func (rs *regState) resetX() { rs.x = regInfo{val: -1} }
func (rs *regState) resetY() { rs.y = regInfo{val: -1} }

// loadImm records that a register was loaded with an immediate value
func (rs *regState) loadImm(reg string, val int) {
	r := regInfo{val: val}
	switch reg {
	case "a":
		rs.a = r
	case "x":
		rs.x = r
	case "y":
		rs.y = r
	}
}

// loadMem records that a register was loaded from a memory location
func (rs *regState) loadMem(reg, varName string) {
	r := regInfo{val: -1, src: varName}
	switch reg {
	case "a":
		rs.a = r
	case "x":
		rs.x = r
	case "y":
		rs.y = r
	}
}

// storeInvalidate marks that a memory location was written.
// If any register's src matches, that register becomes "unknown" (value gone).
func (rs *regState) storeInvalidate(varName string) {
	if rs.a.src == varName {
		rs.a.src = ""
		rs.a.val = -1
	}
	if rs.x.src == varName {
		rs.x.src = ""
		rs.x.val = -1
	}
	if rs.y.src == varName {
		rs.y.src = ""
		rs.y.val = -1
	}
}

// isChainBreaker returns true for instructions that reset all register tracking
func isChainBreaker(opcode string) bool {
	switch opcode {
	case "jmp", "jsr", "rts", "rti", "brk":
		return true
	}
	return false
}

// updateRegState updates the register state machine for a code line
func updateRegState(line asmLine, rs *regState) {
	if !line.isCode {
		return
	}

	switch line.opcode {
	case "lda":
		if strings.HasPrefix(line.operand, "#") {
			if val := parseHexOrDec(line.operand[1:]); val >= 0 {
				rs.loadImm("a", val)
			} else {
				rs.resetA()
			}
		} else {
			rs.loadMem("a", line.operand)
		}
	case "ldx":
		if strings.HasPrefix(line.operand, "#") {
			if val := parseHexOrDec(line.operand[1:]); val >= 0 {
				rs.loadImm("x", val)
			} else {
				rs.resetX()
			}
		} else {
			rs.loadMem("x", line.operand)
		}
	case "ldy":
		if strings.HasPrefix(line.operand, "#") {
			if val := parseHexOrDec(line.operand[1:]); val >= 0 {
				rs.loadImm("y", val)
			} else {
				rs.resetY()
			}
		} else {
			rs.loadMem("y", line.operand)
		}
	case "sta":
		rs.storeInvalidate(line.operand)
	case "stx":
		rs.storeInvalidate(line.operand)
	case "sty":
		rs.storeInvalidate(line.operand)
	case "inc", "dec":
		rs.storeInvalidate(line.operand)
	case "tax":
		rs.x = regInfo{val: -1, src: rs.a.src}
	case "tay":
		rs.y = regInfo{val: -1, src: rs.a.src}
	case "txa":
		rs.a = regInfo{val: -1, src: rs.x.src}
	case "tya":
		rs.a = regInfo{val: -1, src: rs.y.src}
	case "inx", "dex":
		rs.resetX()
	case "iny", "dey":
		rs.resetY()
	case "asl", "lsr", "rol", "ror", "adc", "sbc", "and", "ora", "eor", "pla":
		rs.resetA()
	case "plp":
		rs.reset()
	case "tsx":
		rs.resetX()
	case "txs":
		// SP doesn't affect register tracking
	}
}