import json
import subprocess
import tempfile
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import pandocfilters as pf  # type: ignore

from ssb_pubmd.adapters.storage import LocalFileStorage
from ssb_pubmd.core.document_processor.models import BlockElement
from ssb_pubmd.core.document_processor.models import PandocAst
from ssb_pubmd.core.document_processor.models import PandocElement
from ssb_pubmd.core.document_processor.quarto import quarto_markdown_to_pandoc_ast


@dataclass
class Document:
    """Represents a generic document, backed by the pandoc AST.

    Attributes:
        ast: The pandoc AST of the document, as a dictionary.
        element_index: A mapping from element id's to their position in the AST.
        element_metadata: A mapping from element id's to their metadata.

            * Element metadata will typically be created programmatically in a notebook document.
            * When executing the notebook, metadata is stored in a temporary filed and retrieved after execution.

    Examples:
        ```py
        ast: PandocAst = {
            "pandoc-api-version": [1, 23, 1],
            "meta": {},
            "blocks": [
                {"t": "Div", "c": [["my-highchart", ["ssb"], [["title", "My highchart"]]], []]}
            ],
        }
        document = Document(ast=ast, element_index={}, element_metadata={})
        print(document.extract_html())
        #<div id="my-highchart" class="ssb" title="My highchart"></div>
        ```

    Reference:
        - The best "reference" of the pandoc AST is to run `pandoc FILE -t json`, where FILE is a minimal example document (e.g. in Markdown or html).
        - https://github.com/jgm/pandocfilters has some examples of how to work with the AST.
        - Note that no formal specification exists.
    """

    ast: PandocAst
    element_index: dict[str, int]
    element_metadata: dict[str, dict[str, Any]]

    @classmethod
    def from_quarto_markdown(cls, file_path: str) -> "Document":
        """Creates a document from a .qmd file."""
        with tempfile.NamedTemporaryFile(
            delete=True, suffix=".json"
        ) as temporary_storage_file:
            ast = quarto_markdown_to_pandoc_ast(file_path, temporary_storage_file.name)
            element_index = cls._generate_element_index(ast, target_class="ssb")
            storage = LocalFileStorage(Path(temporary_storage_file.name))
            element_metadata = {key: storage.get(key) for key in element_index.keys()}

            return Document(
                ast=ast,
                element_index=element_index,
                element_metadata=element_metadata,
            )

    def extract_metadata(self, target_key: str) -> dict[str, Any]:
        """Extracts the document metadata stored under a given target key."""

        def meta_to_dict(meta: Any) -> Any:
            t, c = meta.get("t"), meta.get("c")
            if t == "MetaMap":
                return {k: meta_to_dict(v) for k, v in c.items()}
            elif t == "MetaList":
                return [meta_to_dict(v) for v in c]
            else:
                return pf.stringify(c)

        return meta_to_dict(self.ast["meta"][target_key])  # type: ignore

    def extract_html(self) -> str:
        """Extracts the whole document converted to html."""
        return self._document_to_html(self.ast)

    def extract_elements(self, target_class: str) -> Iterator[BlockElement]:
        """Extracts the top-level div elements of a given target class."""
        for id_, i in self.element_index.items():
            metadata = self.element_metadata[id_]
            element = self.ast["blocks"][i]
            inner_blocks: list[PandocElement] = element["c"][1]
            inner_html = self._blocks_to_html(inner_blocks) if inner_blocks else ""
            yield BlockElement(id_, metadata, inner_html)

    def replace_element(self, id_: str, new_html: str) -> None:
        """Replaces the element of a given id with a new html fragment.

        If an element index is not created, or if the given id doesn't exist, it will raise a KeyError.
        """
        i = self.element_index[id_]
        self.ast["blocks"][i] = {
            "t": "RawBlock",
            "c": ["html", new_html],
        }

    @classmethod
    def _blocks_to_html(cls, blocks: list[PandocElement]) -> str:
        document: PandocAst = {
            "pandoc-api-version": [1, 23, 1],
            "meta": {},
            "blocks": blocks,
        }
        return cls._document_to_html(document)

    @classmethod
    def _document_to_html(cls, document: PandocAst) -> str:
        result = subprocess.run(
            ["quarto", "pandoc", "-f", "json", "-t", "html"],
            input=json.dumps(document),
            text=True,
            capture_output=True,
            check=True,
        )
        html = result.stdout
        return html

    @staticmethod
    def _generate_element_index(ast: PandocAst, target_class: str) -> dict[str, int]:
        """Creates an index of all the top-level div elements of a given target class.

        Returs:
            A mapping between element id's and their index in the document.
        """
        index = {}
        for i, element in enumerate(ast["blocks"]):
            if element["t"] != "Div":
                continue

            id_: str = element["c"][0][0]
            if not id_:
                continue

            classes: list[str] = element["c"][0][1]
            if target_class not in classes:
                continue

            index[id_] = i

        return index
