from collections import defaultdict
from dataclasses import dataclass, field
from typing import Dict, List, Tuple, Union

from gbso.cpu.regs import R8

Loc = Union[R8, int]


@dataclass
class CPUState:
    carry: int = 0
    sp: int = 0
    reg8: Dict[R8, int] = field(default_factory=lambda: defaultdict(lambda: 0))
    memory: Dict[int, int] = field(default_factory=lambda: defaultdict(lambda: 0))

    def copy(self) -> "CPUState":
        return CPUState(
            carry=self.carry,
            sp=self.sp,
            reg8=self.reg8.copy(),
            memory=self.memory.copy(),
        )

    def with_vals(self, *init_vals: List[Tuple[Loc, int]]) -> "CPUState":
        cpu = self.copy()
        for loc, val in init_vals:
            if type(loc) == R8:
                cpu.reg8[loc] = val
            else:
                cpu.memory[loc] = val
        return cpu