from contextlib import suppress
from functools import wraps
from pprint import pformat
from typing import Callable
from typing import Dict
from typing import List
from typing import Literal
from typing import Optional
from typing import Sequence
from typing import Tuple
from typing import Type
from typing import TypeVar
from typing import Union
from typing import cast
from typing import overload

from pytezos.michelson.forge import unforge_address
from pytezos.michelson.forge import unforge_chain_id
from pytezos.michelson.forge import unforge_micheline
from pytezos.michelson.forge import unforge_public_key
from pytezos.michelson.forge import unforge_signature
from pytezos.michelson.format import micheline_to_michelson


class MichelsonRuntimeError(Exception):
    def format_stdout(self):
        offset, instruction = next(
            ((len(self.args) - i, arg) for i, arg in enumerate(reversed(self.args)) if str(arg).isupper()), (0, 'ERROR')
        )
        message = ' -> '.join(self.args[offset:])
        return f'{instruction}: {message}'


def catch(prim, func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            if not e.args:
                e.args = (type(e).__name__,)
            if prim:
                e.args = (prim, *e.args)
            raise MichelsonRuntimeError(*e.args) from e

    return wrapper


def try_catch(prim):
    def _catch(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            try:
                return func(*args, **kwargs)
            except Exception as e:
                raise MichelsonRuntimeError(prim, *e.args) from e

        return wrapper

    return _catch


class ErrorTrace(type):
    def __new__(mcs, name, bases, attrs, **kwargs):
        wrapped_attrs = {}
        for attr_name, attr in attrs.items():
            prim = kwargs.get('prim')
            if type(attr) in [classmethod, staticmethod] and not attr_name.startswith('_'):
                attr = type(attr)(catch(prim, attr.__func__))
            elif callable(attr) and not attr_name.startswith('_'):
                attr = catch(prim, attr)
            wrapped_attrs[attr_name] = attr
        return type.__new__(mcs, name, bases, wrapped_attrs, **kwargs)


def parse_micheline_prim(prim_expr) -> Tuple[str, list, list]:
    assert isinstance(prim_expr, dict), f'expected dict, got {pformat(prim_expr)} (instr_expr)'
    prim = prim_expr.get('prim')
    assert prim is not None, f'prim field is absent'
    args = prim_expr.get('args', [])
    assert isinstance(args, list), f'{prim}: expected list of args, got {pformat(args)} (args)'
    annots = prim_expr.get('annots', [])
    assert isinstance(annots, list), f'{prim}: expected list of annots, got {pformat(annots)} (annots)'
    return prim, args, annots


def parse_micheline_value(val_expr, handlers: Dict[Tuple[str, int], Callable]):
    assert isinstance(val_expr, dict), f'expected dict, got {pformat(val_expr)} (val_expr)'
    prim, args = val_expr.get('prim'), val_expr.get('args', [])
    expected = ' or '.join(map(lambda x: f'{x[0]} ({x[1]} args)', handlers))
    assert (prim, len(args)) in handlers, f'expected {expected}, got {prim} ({len(args)} args)'
    handler = handlers[(prim, len(args))]  # type: ignore
    return handler(args)


def parse_micheline_literal(val_expr, handlers: Dict[str, Callable]):
    assert isinstance(val_expr, dict), f'expected dict, got {pformat(val_expr)}'
    try:
        core_type, value = next((k, v) for k, v in val_expr.items() if k[0] != '_' and k != 'annots')
    except StopIteration as e:
        raise Exception(f'Can\'t parse literal `{val_expr}`') from e
    expected = ' or '.join(map(lambda x: f'`{x}`', handlers))
    if core_type not in handlers:
        raise Exception(f'Expected one of {expected}, got {core_type}')
    handler = handlers[core_type]
    return handler(value)


def micheline_value_to_python_object(val_expr):
    if isinstance(val_expr, dict):
        if len(val_expr) == 1:
            prim = next(iter(val_expr))
            if prim == 'string':
                return val_expr[prim]
            elif prim == 'int':
                return int(val_expr[prim])
            elif prim == 'bytes':
                return blind_unpack(bytes.fromhex(val_expr[prim]))
            elif prim == 'bool':
                return val_expr[prim] == 'True'
        elif val_expr.get('prim'):
            prim = val_expr['prim']
            if prim == 'Pair':
                return tuple(micheline_value_to_python_object(x) for x in val_expr['args'])
    return micheline_to_michelson(val_expr)


def blind_unpack(data: bytes):
    with suppress(ValueError):
        return unforge_chain_id(data)
    with suppress(ValueError, KeyError):
        return unforge_address(data)
    with suppress(ValueError, KeyError):
        return unforge_public_key(data)
    with suppress(ValueError):
        return unforge_signature(data)

    if len(data) > 0 and data.startswith(b'\x05'):
        with suppress(ValueError, AssertionError):
            res = unforge_micheline(data[1:])
            return micheline_value_to_python_object(res)

    with suppress(ValueError):
        return data.decode()

    return data


class Micheline(metaclass=ErrorTrace):
    prim: Optional[str] = None
    args: List[Type['Micheline']] = []
    literal: Optional[Union[int, str, bytes]] = None
    classes: Dict[Tuple[str, Optional[int]], Type['Micheline']] = {}

    @classmethod
    def __init_subclass__(cls, prim: Optional[str] = None, args_len: Optional[int] = 0, **kwargs):
        super().__init_subclass__(**kwargs)  # type: ignore
        if prim is not None:
            assert (prim, args_len) not in cls.classes, f'duplicate key {prim} ({args_len} args)'
            cls.classes[(prim, args_len)] = cls
            cls.prim = prim

    def __str__(self):
        raise AssertionError('__str__ has to be explicitly defined')

    def __repr__(self):
        raise AssertionError('__repr__ has to be explicitly defined')

    @staticmethod
    def match(expr) -> Type['Micheline']:
        if isinstance(expr, list):
            args = [Micheline.match(arg) for arg in expr]
            return MichelineSequence.create_type(args=args)
        elif isinstance(expr, dict):
            if expr.get('prim'):
                prim, args, annots = parse_micheline_prim(expr)

                # FIXME: We need entrypoint to be an argument
                if prim == 'RUN':
                    if annots:
                        args = [{'string': annots[0][1:]}] + args  # type: ignore
                        annots = []
                    else:
                        args = [{'string': 'default'}] + args  # type: ignore

                args_len = len(args)
                if (prim, args_len) not in Micheline.classes:
                    args_len = None  # type: ignore
                assert (prim, args_len) in Micheline.classes, f'unregistered primitive {prim} ({args_len} args)'
                cls = Micheline.classes[prim, args_len]
                try:
                    return cls.create_type(args=list(map(Micheline.match, args)), annots=annots)
                except Exception as e:
                    raise MichelsonRuntimeError(cls.prim, *e.args) from e
            else:
                literal = parse_micheline_literal(
                    expr,
                    {
                        'int': int,
                        'string': str,
                        'bytes': bytes.fromhex,
                    },
                )
                return MichelineLiteral.create(literal=literal)
        else:
            raise MichelsonRuntimeError(f'malformed expression `{expr}`')

    @classmethod
    def create_type(
        cls,
        args: List[Type['Micheline']],
        annots: Optional[list] = None,
        **kwargs,
    ) -> Type['Micheline']:
        res = type(cls.__name__, (cls,), dict(args=args, **kwargs))
        return cast(Type['MichelsonPrimitive'], res)  # type: ignore

    @classmethod
    def assert_type_equal(cls, other: Type['Micheline'], path='', message=''):
        comment = f' [{message}]' if message else ''
        assert cls.prim == other.prim, f'expected {other.prim}, got {cls.prim} at `{path}`{comment}'
        assert len(cls.args) == len(
            other.args
        ), f'expected {len(other.args)} args, got {len(cls.args)} at `{path}`{comment}'
        for i, arg in enumerate(other.args):
            cls.args[i].assert_type_equal(arg, path=f'{path}/{i}', message=message)
        assert (
            cls.literal == other.literal
        ), f'expected {other.literal}, got {cls.literal} at `{path}`{comment}'  # type: ignore

    @classmethod
    def assert_type_in(cls, *others: Type['Micheline'], message=''):
        comment = f' [{message}]' if message else ''
        expected = [ty.prim for ty in others]
        assert any(issubclass(cls, ty) for ty in others), f'expected one of {expected}, got {cls.prim}{comment}'

    @classmethod
    def as_micheline_expr(cls) -> dict:
        args = [arg.as_micheline_expr() for arg in cls.args]
        expr = {'prim': cls.prim, 'args': args}
        return {k: v for k, v in expr.items() if v}

    @classmethod
    def execute(cls, stack, stdout, context) -> 'Micheline':
        raise AssertionError(f'`execute` has to be explicitly defined ({cls.prim})')


MichelineT = TypeVar('MichelineT', bound=Micheline)


class MichelineSequence(Micheline):
    def __init__(self, items: List[Micheline]):
        super(MichelineSequence, self).__init__()
        self.items = items

    @classmethod
    def as_micheline_expr(cls) -> list:  # type: ignore
        return [arg.as_micheline_expr() for arg in cls.args]

    @classmethod
    def execute(cls, stack, stdout, context) -> Micheline:
        return cls([arg.execute(stack, stdout, context) for arg in cls.args])


class GlobalConstant(Micheline, prim='constant', args_len=1):
    @classmethod
    def create_type(
        cls,
        args: List[Type['Micheline']],
        annots: Optional[list] = None,
        **kwargs,
    ) -> Type['Micheline']:
        raise RuntimeError('Please, register global constants in advance using context helpers')


class MichelineLiteral(Micheline):
    @classmethod
    def create(cls, literal: Union[int, str, bytes]):
        return cls.create_type(args=[], annots=[], literal=literal)

    @classmethod
    def as_micheline_expr(cls) -> dict:
        if isinstance(cls.literal, int):
            return {'int': str(cls.literal)}
        elif isinstance(cls.literal, str):
            return {'string': cls.literal}
        elif isinstance(cls.literal, bytes):
            return {'bytes': cls.literal.hex()}
        else:
            raise AssertionError(f'unexpected value `{cls.literal}`')

    @classmethod
    def get_int(cls) -> int:
        if not isinstance(cls.literal, int):
            raise TypeError(f'Expected int, got {cls.literal}')  # type: ignore
        return cls.literal

    @classmethod
    def get_string(cls) -> str:
        if not isinstance(cls.literal, str):
            raise TypeError(f'Expected str, got {cls.literal}')  # type: ignore
        return cls.literal

    @classmethod
    def get_bytes(cls) -> bytes:
        if not isinstance(cls.literal, bytes):
            raise TypeError(f'Expected bytes, got {cls.literal}')  # type: ignore
        return cls.literal


def validate_sections(sequence: Type[MichelineSequence], required_sections: Sequence[str]) -> None:
    if len(sequence.args) < len(required_sections):
        raise Exception(f'expected at least {len(required_sections)} sections, got {len(sequence.args)}')
    sequence_sections = {arg.prim for arg in sequence.args}
    if set(required_sections).intersection(sequence_sections) == required_sections:
        raise Exception(f'Unknown sections {sequence_sections}, required: {required_sections}')


@overload
def get_script_section(
    sequence: Type[MichelineSequence],
    cls: None,
    name: str,
    required: Literal[True],
) -> MichelineT:
    ...


@overload
def get_script_section(
    sequence: Type[MichelineSequence],
    cls: None,
    name: str,
    required: Literal[False],
) -> Optional[MichelineT]:
    ...


@overload
def get_script_section(
    sequence: Type[MichelineSequence],
    cls: Type[MichelineT],
    name: None,
    required: Literal[True],
) -> MichelineT:
    ...


@overload
def get_script_section(
    sequence: Type[MichelineSequence],
    cls: Type[MichelineT],
    name: None,
    required: Literal[False],
) -> Optional[MichelineT]:
    ...


def get_script_section(
    sequence,
    cls=None,
    name=None,
    required=False,
):
    try:
        if cls:
            return next(arg for arg in sequence.args if issubclass(arg, cls))
        elif name:
            return next(arg for arg in sequence['code'] if arg['prim'] == name)
        else:
            raise Exception('Either `cls` or `name` must be specified')
    except StopIteration:
        if required:
            raise
    return None


@overload
def get_script_sections(
    sequence: Type[MichelineSequence],
    cls: None,
    name: str,
) -> List[MichelineT]:
    ...


@overload
def get_script_sections(
    sequence: Type[MichelineSequence],
    cls: Type[MichelineT],
    name: None,
) -> List[MichelineT]:
    ...


def get_script_sections(sequence, cls=None, name=None):
    if cls:
        return [arg for arg in sequence.args if issubclass(arg, cls)]
    elif name:
        return [arg for arg in sequence['code'] if arg['prim'] == name]
    else:
        raise Exception('Either `cls` or `name` must be specified')
