from typing import Generator
from typing import List
from typing import Optional
from typing import Tuple
from typing import Type
from typing import Union
from typing import cast

from pytezos.context.abstract import AbstractContext
from pytezos.michelson.micheline import Micheline
from pytezos.michelson.micheline import parse_micheline_value
from pytezos.michelson.types.adt import ADTMixin
from pytezos.michelson.types.adt import Nested
from pytezos.michelson.types.adt import wrap_or
from pytezos.michelson.types.base import MichelsonType
from pytezos.michelson.types.base import Undefined
from pytezos.michelson.types.base import undefined
from pytezos.michelson.types.core import Unit
from pytezos.michelson.types.core import UnitType


class LeftLiteral(Micheline, prim='Left', args_len=1):
    pass


class RightLiteral(Micheline, prim='Right', args_len=1):
    pass


class OrType(MichelsonType, ADTMixin, prim='or', args_len=2):
    is_enum: bool

    def __init__(self, items: Tuple[Union[undefined, MichelsonType], ...]):
        super(OrType, self).__init__()
        self.items = items

    def __repr__(self):
        return f'({" + ".join(map(repr, self.items))})'

    def __iter__(self) -> Generator[Optional[MichelsonType], None, None]:
        yield from self.items  # type: ignore

    def __eq__(self, other):  # type: ignore
        if not isinstance(other, OrType):
            return False
        return all(item == other.items[i] for i, item in enumerate(self.items))

    def __lt__(self, other: 'OrType'):  # type: ignore
        if self.is_left() and other.is_right():
            return True
        elif self.is_left() and other.is_left():
            return self.items[0] < other.items[0]  # type: ignore
        elif self.is_right() and other.is_right():
            return self.items[1] < other.items[1]  # type: ignore
        else:
            return False

    def __hash__(self):
        return hash(self.items)

    def is_left(self) -> bool:
        return isinstance(self.items[0], MichelsonType)

    def is_right(self) -> bool:
        return isinstance(self.items[1], MichelsonType)

    def resolve(self) -> MichelsonType:
        return next(item for item in self if isinstance(item, MichelsonType))

    @staticmethod
    def from_left(left: MichelsonType, right_type: Type[MichelsonType]):
        cls = OrType.create_type(args=[type(left), right_type])
        return cls((left, Undefined))

    @staticmethod
    def from_right(right: MichelsonType, left_type: Type[MichelsonType]):
        cls = OrType.create_type(args=[left_type, type(right)])
        return cls((Undefined, right))

    @classmethod
    def iter_type_args(cls, entrypoints=False, path='') -> Generator[Tuple[str, Type[MichelsonType]], None, None]:
        for i, arg in enumerate(cls.args):
            if issubclass(arg, OrType):
                if entrypoints and arg.field_name:
                    yield path + str(i), arg
                yield from arg.iter_type_args(entrypoints=entrypoints, path=path + str(i))  # type: ignore
            elif entrypoints is False or arg.field_name:
                yield path + str(i), arg

    @classmethod
    def create_type(
        cls,
        args: List[Type['Micheline']],
        annots: Optional[list] = None,
        **kwargs,
    ) -> Type['OrType']:
        def all_units(arguments: List[Type['Micheline']]):
            for arg in arguments:
                if issubclass(arg, OrType):
                    if not all_units(arg.args):
                        return False
                elif not issubclass(arg, UnitType):
                    return False
            return True

        is_enum = all_units(args)
        res = super(OrType, cls).create_type(args=args, annots=annots, is_enum=is_enum, **kwargs)
        return cast(Type['OrType'], res)

    @classmethod
    def generate_pydoc(cls, definitions: list, inferred_name=None, comparable=False):
        name = cls.field_name or cls.type_name or inferred_name or f'{cls.prim}_{len(definitions)}'
        flat_args = cls.get_flat_args(infer_names=True)
        assert isinstance(flat_args, dict), f'sum type has to be named (in the scope of PyTezos)'
        if cls.is_enum:
            doc = ' || '.join(flat_args.keys())
        else:
            variants = [
                (entrypoint, arg.generate_pydoc(definitions, inferred_name=entrypoint))
                for entrypoint, arg in flat_args.items()
            ]
            if comparable:
                doc = ' ||\n\t'.join(f'( "{entrypoint}", {arg_doc} )' for entrypoint, arg_doc in variants)
            else:
                doc = ' ||\n\t'.join(f'{{ "{entrypoint}": {arg_doc} }}' for entrypoint, arg_doc in variants)
        definitions.insert(0, (name, doc))
        return f'${name}'

    @classmethod
    def dummy(cls, context: AbstractContext):
        return cls((cls.args[0].dummy(context), Undefined))

    @classmethod
    def from_micheline_value(cls, val_expr) -> 'OrType':
        value = parse_micheline_value(
            val_expr,
            {
                ('Left', 1): lambda x: (cls.args[0].from_micheline_value(x[0]), Undefined),
                ('Right', 1): lambda x: (Undefined, cls.args[1].from_micheline_value(x[0])),
            },
        )
        return cls(value)

    @classmethod
    def from_python_object(cls, py_obj) -> 'OrType':
        if isinstance(py_obj, list):
            py_obj = tuple(py_obj)
        elif isinstance(py_obj, str):
            assert cls.is_enum, 'string values allowed for enums only'
            py_obj = {py_obj: Unit}
        elif isinstance(py_obj, tuple):
            assert len(py_obj) == 2, f'expected `(entrypoint, value)`, got {py_obj}'
            py_obj = {py_obj[0]: py_obj[1]}

        if isinstance(py_obj, dict):
            assert len(py_obj) == 1, f'single key expected, got {len(py_obj)}'
            entrypoint = next(iter(py_obj))
            _, key_to_path, _ = cls.get_type_layout(infer_names=True)
            assert key_to_path, f'sum type has to be named (in the scope of PyTezos)'
            return cls.from_python_object(wrap_or(py_obj[entrypoint], key_to_path[entrypoint]))
        elif isinstance(py_obj, Nested):
            value = tuple(
                Undefined if py_obj[i] is Undefined else cls.args[i].from_python_object(py_obj[i]) for i in [0, 1]
            )
            return cls(value)
        else:
            raise AssertionError(f'expected list, tuple, or dict, got `{py_obj}`')

    def iter_values(self, path='') -> Generator[Tuple[str, MichelsonType], None, None]:
        for i, arg in enumerate(self.items):
            if isinstance(arg, OrType):
                yield from arg.iter_values(path + str(i))
            elif isinstance(arg, MichelsonType):
                yield path + str(i), arg
            else:
                assert arg == Undefined, f'expected Michelson type or undefined, got {arg}'

    def to_literal(self) -> Type[Micheline]:
        if self.is_left():
            return LeftLiteral.create_type(args=[self.items[0].to_literal()])  # type: ignore
        else:
            return RightLiteral.create_type(args=[self.items[1].to_literal()])  # type: ignore

    def to_micheline_value(self, mode='readable', lazy_diff=False):
        for i, prim in enumerate(['Left', 'Right']):
            if isinstance(self.items[i], MichelsonType):
                return {'prim': prim, 'args': [self.items[i].to_micheline_value(mode=mode, lazy_diff=lazy_diff)]}
        raise AssertionError(f'unexpected value {self.items}')

    def to_python_object(self, try_unpack=False, lazy_diff=False, comparable=False) -> Union[tuple, dict]:
        flat_values = self.get_flat_values(infer_names=True)
        assert (
            isinstance(flat_values, dict) and len(flat_values) == 1
        ), f'sum type has to be named (in the scope of PyTezos)'
        entrypoint = next(iter(flat_values))
        if self.is_enum:
            return entrypoint  # type: ignore
        else:
            py_obj = flat_values[entrypoint].to_python_object(
                try_unpack=try_unpack,
                lazy_diff=lazy_diff,
                comparable=comparable,
            )
            return (entrypoint, py_obj) if comparable else {entrypoint: py_obj}

    def merge_lazy_diff(self, lazy_diff: List[dict]) -> 'OrType':
        items = tuple(
            item.merge_lazy_diff(lazy_diff) if isinstance(item, MichelsonType) else item for item in self.items
        )
        return type(self)(items)  # type: ignore

    def aggregate_lazy_diff(self, lazy_diff: List[dict], mode='readable'):
        items = tuple(
            item.aggregate_lazy_diff(lazy_diff, mode=mode) if isinstance(item, MichelsonType) else item
            for item in self.items
        )
        return type(self)(items)  # type: ignore

    def attach_context(self, context: AbstractContext, big_map_copy=False):
        for item in self:
            if isinstance(item, MichelsonType):
                item.attach_context(context, big_map_copy=big_map_copy)

    def __getitem__(self, key: Union[int, str]) -> MichelsonType:
        return self.get_value(key, infer_names=True)
