from dataclasses import replace
from enum import Enum
from random import choice, randint, random, randrange
from typing import List, Tuple, Type, TypeVar, Union

from gbso.cpu.insn import ALL_INSN_CLASSES, Insn, Operand, UNUSED, get_signature_class
from gbso.cpu.regs import R16, R16_WITHOUT_SP, R8
from gbso.program.program import Program


class Mutation(Enum):
    OPCODE = "OPCODE"
    OPERAND = "OPERAND"
    SWAP = "SWAP"
    INSTRUCTION = "INSTRUCTION"
    NONE = "NONE"


def mutate_program(
    prgm: Program,
    prob_opcode: float,
    prob_operand: float,
    prob_swap: float,
    prob_insn: float,
    prob_insn_unused: float,
) -> Program:
    prob_none = max(0, 1 - prob_opcode - prob_operand - prob_swap - prob_insn)
    mutation = sample_probs(
        [
            (Mutation.OPCODE, prob_opcode),
            (Mutation.OPERAND, prob_operand),
            (Mutation.SWAP, prob_swap),
            (Mutation.INSTRUCTION, prob_insn),
            (Mutation.NONE, prob_none),
        ]
    )
    new_prgm = prgm

    if mutation == Mutation.OPCODE:
        new_prgm = mutate_opcode(prgm)
    elif mutation == Mutation.OPERAND:
        new_prgm = mutate_operand(prgm)
    elif mutation == Mutation.SWAP:
        new_prgm = mutate_swap(prgm)
    elif mutation == Mutation.INSTRUCTION:
        new_prgm = mutate_insn(prgm, prob_insn_unused)

    return new_prgm


def mutate_opcode(prgm: Program) -> Program:
    i = randrange(len(prgm.insns))
    next_insn_cls = choice(get_signature_class(prgm.insns[i]))

    insns = prgm.insns.copy()
    insns[i] = next_insn_cls(*vars(prgm.insns[i]).values())

    return Program(insns=insns)


# God, this is ugly... a better approach would be to store operands as a tuple
# instead of individual named fields.
def mutate_operand(prgm: Program) -> Program:
    insn_index = randrange(len(prgm.insns))
    insn = prgm.insns[insn_index]

    if len(insn.signature()) == 0:
        return prgm

    operand_index = randrange(len(insn.signature()))
    operand_type = insn.signature()[operand_index]
    operand_name = list(vars(insn).keys())[operand_index]

    insns = prgm.insns.copy()
    insns[insn_index] = replace(
        insn, **{operand_name: create_random_operand(operand_type)}
    )

    return Program(insns=insns)


def mutate_swap(prgm: Program) -> Program:
    i = randrange(len(prgm.insns))
    j = randrange(len(prgm.insns))

    insns = prgm.insns.copy()
    insn = insns[i]
    insns[i] = insns[j]
    insns[j] = insn

    return Program(insns=insns)


def mutate_insn(prgm: Program, prob_unused: float) -> Program:
    unused = random()
    next_insn: Insn = UNUSED()
    if unused > prob_unused:
        cls = choice(ALL_INSN_CLASSES)
        next_insn = create_random_insn(cls)

    insns = prgm.insns.copy()
    insn_index = randrange(len(prgm.insns))
    insns[insn_index] = next_insn

    return Program(insns=insns)


A = TypeVar("A")


# NOTE: List must be non-empty and probabilities must sum to 1
def sample_probs(probs: List[Tuple[A, float]]) -> A:
    r = random()
    cum_prob = 0.0

    for val, prob in probs[:-1]:
        cum_prob += prob
        if r < cum_prob:
            return val

    val, _ = probs[-1]
    return val


def create_random_insn(insn_cls: Type[Insn]) -> Insn:
    args = [create_random_operand(op) for op in insn_cls.signature()]
    return insn_cls(*args)  # type: ignore


def create_random_operand(op: Operand) -> Union[int, R8, R16]:
    value: Union[int, R8, R16]

    if op == Operand.R8:
        value = choice(list(R8))
    elif op == Operand.R16:
        value = choice(list(R16))
    elif op == Operand.R16_NO_SP:
        value = choice(R16_WITHOUT_SP)
    elif op == Operand.IMM3:
        value = randrange(8)
    elif op == Operand.IMM8:
        value = randrange(256)
    elif op == Operand.IMM16:
        value = randrange(256)
    elif op == Operand.SIMM8:
        value = randint(-128, 127)

    return value