from typing import List
from typing import Tuple
from typing import Union
from typing import cast

from pymavryk.context.abstract import AbstractContext
from pymavryk.michelson.instructions.base import MichelsonInstruction
from pymavryk.michelson.instructions.base import format_stdout
from pymavryk.michelson.stack import MichelsonStack
from pymavryk.michelson.types import BigMapType
from pymavryk.michelson.types import BoolType
from pymavryk.michelson.types import ListType
from pymavryk.michelson.types import MapType
from pymavryk.michelson.types import MichelsonType
from pymavryk.michelson.types import OptionType
from pymavryk.michelson.types import SetType


class ConsInstruction(MichelsonInstruction, prim='CONS'):
    @classmethod
    def execute(cls, stack: MichelsonStack, stdout: List[str], context: AbstractContext):
        elt, lst = cast(Tuple[MichelsonType, ListType], stack.pop2())
        lst.assert_type_in(ListType)
        res = lst.prepend(elt)
        stack.push(res)
        stdout.append(format_stdout(cls.prim, [elt, lst], [res]))  # type: ignore
        return cls(stack_items_added=1)


class NilInstruction(MichelsonInstruction, prim='NIL', args_len=1):
    @classmethod
    def execute(cls, stack: MichelsonStack, stdout: List[str], context: AbstractContext):
        res = ListType.empty(cls.args[0])  # type: ignore
        stack.push(res)
        stdout.append(format_stdout(cls.prim, [], [res]))  # type: ignore
        return cls(stack_items_added=1)


class EmptyBigMapInstruction(MichelsonInstruction, prim='EMPTY_BIG_MAP', args_len=2):
    @classmethod
    def execute(cls, stack: MichelsonStack, stdout: List[str], context: AbstractContext):
        res = BigMapType.empty(key_type=cls.args[0], val_type=cls.args[1])  # type: ignore
        res.attach_context(context)
        stack.push(res)
        stdout.append(format_stdout(cls.prim, [], [res]))  # type: ignore
        return cls(stack_items_added=1)


class EmptyMapInstruction(MichelsonInstruction, prim='EMPTY_MAP', args_len=2):
    @classmethod
    def execute(cls, stack: MichelsonStack, stdout: List[str], context: AbstractContext):
        res = MapType.empty(key_type=cls.args[0], val_type=cls.args[1])  # type: ignore
        stack.push(res)
        stdout.append(format_stdout(cls.prim, [], [res]))  # type: ignore
        return cls(stack_items_added=1)


class EmptySetInstruction(MichelsonInstruction, prim='EMPTY_SET', args_len=1):
    @classmethod
    def execute(cls, stack: MichelsonStack, stdout: List[str], context: AbstractContext):
        res = SetType.empty(item_type=cls.args[0])  # type: ignore
        stack.push(res)
        stdout.append(format_stdout(cls.prim, [], [res]))  # type: ignore
        return cls(stack_items_added=1)


class GetInstruction(MichelsonInstruction, prim='GET'):
    @classmethod
    def execute(cls, stack: MichelsonStack, stdout: List[str], context: AbstractContext):
        key, src = cast(Tuple[MichelsonType, Union[MapType, BigMapType]], stack.pop2())
        src.assert_type_in(MapType, BigMapType)
        val = src.get(key, dup=True)
        if val is None:
            res = OptionType.none(src.args[1])
        else:
            res = OptionType.from_some(val)
        stack.push(res)
        stdout.append(format_stdout(cls.prim, [key, src], [res]))  # type: ignore
        return cls(stack_items_added=1)


class GetAndUpdateInstruction(MichelsonInstruction, prim='GET_AND_UPDATE'):
    @classmethod
    def execute(cls, stack: MichelsonStack, stdout: List[str], context: AbstractContext):
        key, val, src = cast(Tuple[MichelsonType, OptionType, Union[MapType, BigMapType]], stack.pop3())
        src.assert_type_in(MapType, BigMapType)
        prev_val, dst = src.update(key, None if val.is_none() else val.get_some())
        res = OptionType.none(src.args[1]) if prev_val is None else OptionType.from_some(prev_val)
        stack.push(dst)
        stack.push(res)
        stdout.append(format_stdout(cls.prim, [key, val, src], [res, dst]))  # type: ignore
        return cls(stack_items_added=2)


class UpdateInstruction(MichelsonInstruction, prim='UPDATE'):
    @classmethod
    def execute(cls, stack: MichelsonStack, stdout: List[str], context: AbstractContext):
        key, val, src = cast(
            Tuple[
                MichelsonType,
                Union[OptionType, BoolType],
                Union[MapType, BigMapType, SetType],
            ],
            stack.pop3(),
        )
        val.assert_type_in(OptionType, BoolType)
        if isinstance(val, BoolType):
            src.assert_type_in(SetType)
            dst = src.add(key) if bool(val) else src.remove(key)  # type: ignore
        else:
            src.assert_type_in(MapType, BigMapType)
            _, dst = src.update(key, None if val.is_none() else val.get_some())  # type: ignore
        stack.push(dst)
        stdout.append(format_stdout(cls.prim, [key, val, src], [dst]))  # type: ignore
        return cls(stack_items_added=1)


class MemInstruction(MichelsonInstruction, prim='MEM'):
    @classmethod
    def execute(cls, stack: MichelsonStack, stdout: List[str], context: AbstractContext):
        key, src = cast(Tuple[MichelsonType, Union[SetType, MapType, BigMapType]], stack.pop2())
        src.assert_type_in(MapType, BigMapType, SetType)
        res = BoolType.from_value(src.contains(key))
        stack.push(res)
        stdout.append(format_stdout(cls.prim, [key, src], [res]))  # type: ignore
        return cls(stack_items_added=1)


class NoneInstruction(MichelsonInstruction, prim='NONE', args_len=1):
    @classmethod
    def execute(cls, stack: MichelsonStack, stdout: List[str], context: AbstractContext):
        res = OptionType.none(cls.args[0])  # type: ignore
        stack.push(res)
        stdout.append(format_stdout(cls.prim, [], [res]))  # type: ignore
        return cls(stack_items_added=1)


class SomeInstruction(MichelsonInstruction, prim='SOME'):
    @classmethod
    def execute(cls, stack: MichelsonStack, stdout: List[str], context: AbstractContext):
        some = stack.pop1()
        res = OptionType.from_some(some)
        stack.push(res)
        stdout.append(format_stdout(cls.prim, [some], [res]))  # type: ignore
        return cls(stack_items_added=1)
