from __future__ import annotations

"""Simple JSON-based storage for neuraparse artefacts.

For v0.1 we use one JSON file per document / graph under a user-provided
output directory. This keeps things transparent and easy to inspect.
"""

from pathlib import Path
import hashlib
import json
from typing import Dict, Tuple

from ..core.models import (
    DocumentGraph,
    DocumentNode,
    DocumentTree,
    GraphEdge,
    GraphNode,
    NodeType,
    RawDocument,
    SourceType,
)


def make_document_id(raw: RawDocument) -> str:
    """Derive a stable, filesystem-safe identifier for a document.

    We hash the original ``raw.id`` and keep the first 16 hex characters.
    The original identifier is always stored in the payload metadata.
    """

    h = hashlib.sha1(raw.id.encode("utf-8")).hexdigest()
    return h[:16]


def _node_to_dict(node: DocumentNode) -> Dict:
    return {
        "id": node.id,
        "node_type": node.node_type.value,
        "text": node.text,
        "metadata": node.metadata,
        "children": [_node_to_dict(c) for c in node.children],
    }


def _node_from_dict(data: Dict) -> DocumentNode:
    return DocumentNode(
        id=data["id"],
        node_type=NodeType(data["node_type"]),
        text=data.get("text"),
        metadata=data.get("metadata", {}),
        children=[_node_from_dict(c) for c in data.get("children", [])],
    )


def save_document(base_dir: str | Path, document_id: str, raw: RawDocument, tree: DocumentTree) -> Path:
    """Persist ``raw`` and ``tree`` as a single JSON file.

    The file is placed under ``<base_dir>/documents/<document_id>.json``.
    """

    base = Path(base_dir)
    docs_dir = base / "documents"
    docs_dir.mkdir(parents=True, exist_ok=True)

    payload = {
        "document_id": document_id,
        "raw": {
            "id": raw.id,
            "source_type": raw.source_type.value,
            "content": raw.content,
            "metadata": raw.metadata,
        },
        "tree": {
            "root": _node_to_dict(tree.root),
            "metadata": tree.metadata,
        },
    }

    path = docs_dir / f"{document_id}.json"
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
    return path


def load_document(base_dir: str | Path, document_id: str) -> Tuple[RawDocument, DocumentTree]:
    """Load a previously saved document bundle."""

    base = Path(base_dir)
    path = base / "documents" / f"{document_id}.json"
    data = json.loads(path.read_text(encoding="utf-8"))

    raw_data = data["raw"]
    tree_data = data["tree"]

    raw = RawDocument(
        id=raw_data["id"],
        source_type=SourceType(raw_data["source_type"]),
        content=raw_data["content"],
        metadata=raw_data.get("metadata", {}),
    )

    root = _node_from_dict(tree_data["root"])
    tree = DocumentTree(root=root, metadata=tree_data.get("metadata", {}))
    return raw, tree


def _graph_to_dict(graph: DocumentGraph) -> Dict:
    return {
        "metadata": graph.metadata,
        "nodes": [
            {
                "id": n.id,
                "node_type": n.node_type.value,
                "label": n.label,
                "metadata": n.metadata,
            }
            for n in graph.nodes
        ],
        "edges": [
            {
                "source": e.source,
                "target": e.target,
                "edge_type": e.edge_type,
                "metadata": e.metadata,
            }
            for e in graph.edges
        ],
    }


def _graph_from_dict(data: Dict) -> DocumentGraph:
    nodes = [
        GraphNode(
            id=n["id"],
            node_type=NodeType(n["node_type"]),
            label=n.get("label"),
            metadata=n.get("metadata", {}),
        )
        for n in data.get("nodes", [])
    ]
    edges = [
        GraphEdge(
            source=e["source"],
            target=e["target"],
            edge_type=e["edge_type"],
            metadata=e.get("metadata", {}),
        )
        for e in data.get("edges", [])
    ]
    return DocumentGraph(nodes=nodes, edges=edges, metadata=data.get("metadata", {}))


def save_graph(base_dir: str | Path, document_id: str, graph: DocumentGraph) -> Path:
    """Persist a :class:`DocumentGraph` to JSON under ``<base_dir>/graphs``."""

    base = Path(base_dir)
    graphs_dir = base / "graphs"
    graphs_dir.mkdir(parents=True, exist_ok=True)

    payload = {
        "document_id": document_id,
        "graph": _graph_to_dict(graph),
    }

    path = graphs_dir / f"{document_id}.json"
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
    return path


def load_graph(base_dir: str | Path, document_id: str) -> DocumentGraph:
    """Load a previously saved :class:`DocumentGraph`."""

    base = Path(base_dir)
    path = base / "graphs" / f"{document_id}.json"
    data = json.loads(path.read_text(encoding="utf-8"))
    return _graph_from_dict(data["graph"])

