from typing import Dict, Optional, Tuple

from pydantic import BaseModel, root_validator
from typing_extensions import Literal

from tdm.future.abstract.datamodel import AbstractNode, TalismanDocument
from tdm.future.datamodel.document import TalismanDocumentFactory
from tdm.json_schema import DocumentMetadataModel
from .directives import DirectivesModel
from .facts import FactsModel
from .links import NodeLinksModel
from .nodes import NodeModel, fill_children, serialize_node


class TalismanDocumentModel(BaseModel):
    VERSION: Literal['1.0'] = ...
    id: str
    main_node: str
    content: Tuple[NodeModel, ...]
    links: NodeLinksModel
    facts: FactsModel
    directives: DirectivesModel

    metadata: Optional[DocumentMetadataModel]  # should be removed in future

    @root_validator(pre=True)
    def _set_version(cls, values):  # noqa N805: pydantic requires cls to be first argument in validators
        if 'VERSION' not in values:
            values['VERSION'] = '1.0'
        return values

    def deserialize(self) -> TalismanDocument:
        id2node, structure = self._collect_nodes()

        return TalismanDocumentFactory.construct(
            content=id2node.values(),
            structure=structure,
            root=self.main_node,
            node_links=self.links.deserialize(id2node),
            facts=self.facts.deserialize(id2node),
            directives=self.directives.deserialize(id2node),
            metadata=self.metadata.to_metadata() if self.metadata is not None else None,
            id_=self.id
        )

    def _collect_nodes(self) -> Tuple[Dict[str, AbstractNode], Dict[str, Tuple[str, ...]]]:
        # try to avoid recursion
        links: Dict[str, Tuple[str, ...]] = {}
        id2node = {}

        for node_model in self.content:
            id2node[node_model.id] = node_model.deserialize({})
            if node_model.children:
                links[node_model.id] = node_model.children
        return id2node, links

    @classmethod
    def serialize(cls, document: TalismanDocument) -> 'TalismanDocumentModel':
        node_models = tuple(serialize_node(node) for node in document.nodes.values())
        fill_children(node_models, document)
        return cls(
            id=document.id,
            main_node=document.main_root.id,
            content=node_models,
            links=NodeLinksModel.serialize(document.node_links),
            facts=FactsModel.serialize(document.facts),
            directives=DirectivesModel.serialize(document.directives),
            metadata=DocumentMetadataModel(**document._metadata) if document._metadata is not None else None
        )
