from collections import defaultdict
from typing import Callable, Dict, List, Tuple

from tdm.abstract.datamodel import AbstractNode, AbstractNodeLink, BaseNodeMetadata, FrozenMarkup
from tdm.datamodel.links import TranslationNodeLink
from tdm.datamodel.mentions import NodeMention
from tdm.datamodel.nodes import ImageNode, ImageNodeMetadata, JSONNode, KeyNode, ListNode, ListNodeMetadata, TableCellNode, \
    TableNode, TableRowNode, TableRowNodeMetadata, TextNode, TextNodeMetadata
from tdm.v0.json_schema.content import NodeType, TreeDocumentContentModel


def _convert_text_node(node: TreeDocumentContentModel) -> TextNode:
    return TextNode(
        content=node.text,
        markup=FrozenMarkup.freeze(node.markup.dict()),
        id=node.id,
        metadata=TextNodeMetadata(hidden=node.metadata.hidden, language=node.metadata.language)
    )


def _covert_key_node(node: TreeDocumentContentModel) -> KeyNode:
    return KeyNode(
        content=node.text,
        markup=FrozenMarkup.freeze(node.markup.dict()),
        id=node.id,
        metadata=TextNodeMetadata(hidden=node.metadata.hidden, language=node.metadata.language)
    )


def _convert_list_node(node: TreeDocumentContentModel) -> ListNode:
    return ListNode(id=node.id, metadata=ListNodeMetadata(hidden=node.metadata.hidden))


def _convert_json_node(node: TreeDocumentContentModel) -> JSONNode:
    return JSONNode(id=node.id, metadata=BaseNodeMetadata(hidden=node.metadata.hidden))


def _convert_table_node(node: TreeDocumentContentModel) -> TableNode:
    return TableNode(id=node.id, metadata=BaseNodeMetadata(hidden=node.metadata.hidden))


def _convert_table_row_node(node: TreeDocumentContentModel) -> TableRowNode:
    return TableRowNode(id=node.id, metadata=TableRowNodeMetadata(hidden=node.metadata.hidden))


def _convert_image_node(node: TreeDocumentContentModel) -> ImageNode:
    return ImageNode(
        content=node.text,
        markup=FrozenMarkup.freeze(node.markup.dict()),
        id=node.id,
        metadata=ImageNodeMetadata(hidden=node.metadata.hidden)
    )


NODE_CONVERTERS: Dict[NodeType, Callable[[TreeDocumentContentModel], AbstractNode]] = {
    NodeType.HEADER: _convert_text_node,
    NodeType.TEXT: _convert_text_node,
    NodeType.LIST: _convert_list_node,
    NodeType.JSON: _convert_json_node,
    NodeType.KEY: _covert_key_node,
    NodeType.TABLE: _convert_table_node,
    NodeType.TABLE_ROW: _convert_table_row_node,
    NodeType.IMAGE: _convert_image_node
}


def build_structure(root: TreeDocumentContentModel) -> Tuple[List[AbstractNode], Dict[str, List[str]], List[AbstractNodeLink]]:
    nodes: list[AbstractNode] = []
    structure: dict[str, list[str]] = defaultdict(list)
    node_links: list[AbstractNodeLink] = []

    to_process = [root]  # old document is single tree, so we can start from root

    for node in to_process:
        # create new node for each document node with same type
        converted = NODE_CONVERTERS[node.metadata.node_type](node)
        nodes.append(converted)

        # create separate node with original text (current node is target, original text is source), type is the same
        if node.metadata.original_text is not None:
            if node.metadata.node_type is not NodeType.KEY:
                raise ValueError
            node_links.append(TranslationNodeLink(
                source=NodeMention(KeyNode(content=node.metadata.original_text)),
                target=NodeMention(converted)
            ))

        # text translations are moved to separate nodes linked by translation link (current node is source, translation is target)
        if node.metadata.text_translations:
            if node.metadata.node_type is not NodeType.TEXT and node.metadata.node_type is not NodeType.HEADER:
                raise ValueError
            node_links.extend(
                TranslationNodeLink(
                    source=NodeMention(converted),
                    target=NodeMention(TextNode(content=text, metadata=TextNodeMetadata(language=lang)))
                ) for lang, text in node.metadata.text_translations.items()
            )

        children = node.nodes
        if children:
            if isinstance(converted, TableRowNode):
                # table row should contain only cell nodes (there are no such nodes in old document)
                for c in children:
                    cell = TableCellNode()
                    nodes.append(cell)
                    structure[converted.id].append(cell.id)
                    structure[cell.id].append(c.id)
            else:
                structure[converted.id].extend(map(lambda n: n.id, children))
            to_process.extend(children)

    return nodes, structure, node_links
