"""
Knowledge serialization — save and load UniversalKnowledge to/from disk.
"""

from __future__ import annotations

import pickle
from pathlib import Path
from typing import Union

from model_osmosis.knowledge.representation import UniversalKnowledge


def save_knowledge(knowledge: UniversalKnowledge, path: Union[str, Path]) -> Path:
    """Serialize a UniversalKnowledge object to disk.

    Args:
        knowledge: The knowledge bundle to save.
        path: File path to write to. Will be created (with parents) if needed.
            Recommended extension: ``.osm``

    Returns:
        The resolved path that was written to.
    """
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "wb") as f:
        pickle.dump(knowledge, f, protocol=pickle.HIGHEST_PROTOCOL)
    return path


def load_knowledge(path: Union[str, Path]) -> UniversalKnowledge:
    """Deserialize a UniversalKnowledge object from disk.

    Args:
        path: File path to read from.

    Returns:
        The loaded UniversalKnowledge object.

    Raises:
        FileNotFoundError: If the file does not exist.
        TypeError: If the loaded object is not a UniversalKnowledge instance.
    """
    path = Path(path)
    if not path.exists():
        raise FileNotFoundError(f"Knowledge file not found: {path}")
    with open(path, "rb") as f:
        obj = pickle.load(f)
    if not isinstance(obj, UniversalKnowledge):
        raise TypeError(
            f"Expected UniversalKnowledge, got {type(obj).__name__}"
        )
    return obj
