from functools import lru_cache
from operator import attrgetter
from typing import Type, Callable, TypeVar, Optional

from rewrite.java import *
from rewrite.remote import SenderContext, ReceiverContext

T = TypeVar('T')

def receive_container(container: Optional[JContainer[T]], _: Optional[str], ctx: ReceiverContext) -> JContainer[T]:
    if container is not None:
        container = container.with_before(ctx.receive_node(container.before, receive_space))
        container = container.padding.with_elements(ctx.receive_nodes(container.padding.elements, receive_right_padded_tree))
        container = container.with_markers(ctx.receive_node(container.markers, ctx.receive_markers))
    else:
        container = JContainer(
            ctx.receive_node(None, receive_space),
            ctx.receive_nodes(None, receive_right_padded_tree),
            ctx.receive_node(None, ctx.receive_markers)
        )
    return container


def send_container(container: JContainer[T], ctx: SenderContext):
    ctx.send_node(container, attrgetter('before'), send_space)
    ctx.send_nodes(container, attrgetter('padding.elements'), send_right_padded, lambda t: t.element.id)
    ctx.send_node(container, attrgetter('markers'), ctx.send_markers)


def send_left_padded(left_padded: JLeftPadded[T], ctx: SenderContext):
    ctx.send_node(left_padded, attrgetter('before'), send_space)
    if isinstance(left_padded.element, (J, Space)):
        ctx.send_node(left_padded, attrgetter('element'), ctx.send_tree)
    else:
        ctx.send_value(left_padded, attrgetter('element'))
    ctx.send_node(left_padded, attrgetter('markers'), ctx.send_markers)


def send_right_padded(right_padded: JRightPadded[T], ctx: SenderContext):
    if isinstance(right_padded.element, J):
        ctx.send_node(right_padded, attrgetter('element'), ctx.send_tree)
    elif isinstance(right_padded.element, Space):
        ctx.send_node(right_padded, attrgetter('element'), send_space)
    else:
        ctx.send_value(right_padded, attrgetter('element'))
    ctx.send_node(right_padded, attrgetter('after'), send_space)
    ctx.send_node(right_padded, attrgetter('markers'), ctx.send_markers)


def receive_space(space: Optional[Space], _: Optional[str], ctx: ReceiverContext) -> Space:
    if space is not None:
        space = space.with_comments(ctx.receive_nodes(space.comments, receive_comment))
        space = space.with_whitespace(ctx.receive_value(space.whitespace, str))
    else:
        space = Space(
            ctx.receive_nodes(None, receive_comment),
            ctx.receive_value(None, str)
        )

    return space


def receive_comment(comment: Optional[Comment], _: Optional[str], ctx: ReceiverContext) -> Comment:
    if comment:
        comment = comment.with_multiline(ctx.receive_value(comment.multiline, bool))
        comment = comment.with_text(ctx.receive_value(comment.text, str))
        comment = comment.with_suffix(ctx.receive_value(comment.suffix, str))
        comment = comment.with_markers(ctx.receive_node(comment.markers, ctx.receive_markers))
    else:
        comment = TextComment(
            ctx.receive_value(None, bool),
            ctx.receive_value(None, str),
            ctx.receive_value(None, str),
            ctx.receive_node(None, ctx.receive_markers)
        )
    return comment


def send_space(space: Space, ctx: SenderContext):
    ctx.send_nodes(space, attrgetter('comments'), send_comment, lambda x: x)
    ctx.send_value(space, attrgetter('whitespace'))


def send_comment(comment: Comment, ctx: SenderContext):
    ctx.send_value(comment, attrgetter('multiline'))
    ctx.send_value(comment, attrgetter('text'))
    ctx.send_value(comment, attrgetter('suffix'))
    ctx.send_node(comment, attrgetter('markers'), ctx.send_markers)


@lru_cache(maxsize=10)
def left_padded_value_receiver(type_: Type) -> Callable[[Optional[JLeftPadded[T]], Optional[str], ReceiverContext], JLeftPadded[T]]:
    def receiver(left_padded: Optional[JLeftPadded[T]], _: Optional[str], ctx: ReceiverContext) -> JLeftPadded[T]:
        if left_padded is not None:
            left_padded = left_padded.with_before(ctx.receive_node(left_padded.before, receive_space))
            left_padded = left_padded.with_element(ctx.receive_value(left_padded.element, type_))
            left_padded = left_padded.with_markers(ctx.receive_node(left_padded.markers, ctx.receive_markers))
        else:
            left_padded = JLeftPadded(
                ctx.receive_node(None, receive_space),
                ctx.receive_value(None, type_),
                ctx.receive_node(None, ctx.receive_markers)
            )
        return left_padded
    return receiver


@lru_cache(maxsize=10)
def left_padded_node_receiver(type_: Type) -> Callable[[Optional[JLeftPadded[T]], Optional[str], ReceiverContext], JLeftPadded[T]]:
    if type_ is Space:
        def space_receiver(left_padded: Optional[JLeftPadded[T]], _: Optional[str], ctx: ReceiverContext) -> JLeftPadded[T]:
            if left_padded is not None:
                left_padded = left_padded.with_before(ctx.receive_node(left_padded.before, receive_space))
                left_padded = left_padded.with_element(ctx.receive_node(left_padded.element, receive_space))
                left_padded = left_padded.with_markers(ctx.receive_node(left_padded.markers, ctx.receive_markers))
            else:
                left_padded = JLeftPadded(
                    ctx.receive_node(None, receive_space),
                    ctx.receive_node(None, receive_space),
                    ctx.receive_node(None, ctx.receive_markers)
                )
            return left_padded
        return space_receiver
    else:
        raise ValueError("Only Space is supported")


def receive_left_padded_tree(left_padded: Optional[JLeftPadded[T]], _: Optional[str], ctx: ReceiverContext) -> JLeftPadded[T]:
    if left_padded is not None:
        left_padded = left_padded.with_before(ctx.receive_node(left_padded.before, receive_space))
        left_padded = left_padded.with_element(ctx.receive_node(left_padded.element, ctx.receive_tree))
        left_padded = left_padded.with_markers(ctx.receive_node(left_padded.markers, ctx.receive_markers))
    else:
        left_padded = JLeftPadded(
            ctx.receive_node(None, receive_space),
            ctx.receive_node(None, ctx.receive_tree),
            ctx.receive_node(None, ctx.receive_markers)
        )
    return left_padded


@lru_cache(maxsize=10)
def right_padded_value_receiver(type_: Type) -> Callable[[Optional[JRightPadded[T]], Optional[str], ReceiverContext], JRightPadded[T]]:
    def receiver(right_padded: Optional[JRightPadded[T]], _: Optional[str], ctx: ReceiverContext) -> JRightPadded[T]:
        if right_padded is not None:
            right_padded = right_padded.with_element(ctx.receive_value(right_padded.element, type_))
            right_padded = right_padded.with_after(ctx.receive_node(right_padded.after, receive_space))
            right_padded = right_padded.with_markers(ctx.receive_node(right_padded.markers, ctx.receive_markers))
        else:
            right_padded = JRightPadded(
                ctx.receive_value(None, type_),
                ctx.receive_node(None, receive_space),
                ctx.receive_node(None, ctx.receive_markers)
            )
        return right_padded
    return receiver


@lru_cache(maxsize=10)
def right_padded_node_receiver(type_: Type) -> Callable[[Optional[JRightPadded[T]], Optional[str], ReceiverContext], JRightPadded[T]]:
    raise ValueError("Not implemented")


def receive_right_padded_tree(right_padded: Optional[JRightPadded[T]], _: Optional[str],
                              ctx: ReceiverContext) -> JRightPadded[T]:
    if right_padded is not None:
        right_padded = right_padded.with_element(ctx.receive_node(right_padded.element, ctx.receive_tree))
        right_padded = right_padded.with_after(ctx.receive_node(right_padded.after, receive_space))
        right_padded = right_padded.with_markers(ctx.receive_node(right_padded.markers, ctx.receive_markers))
    else:
        right_padded = JRightPadded(
            ctx.receive_node(None, ctx.receive_tree),
            ctx.receive_node(None, receive_space),
            ctx.receive_node(None, ctx.receive_markers)
        )
    return right_padded
