"""In-memory navigation of the Diffbot Knowledge Graph ontology.

The ontology is a JSON document describing the Knowledge Graph's entity types,
composite types, enums, and taxonomies. An agent constructing DQL needs it to
look up real field paths and taxonomy values instead of guessing them.

This module is pure and storage-agnostic: build an :class:`Ontology` from
already-parsed data (or from raw JSON / a file path) and query it. How the
ontology document is fetched, and whether or where it is cached, is left
entirely to the caller — the `db` CLI caches it on disk at
``~/.diffbot/ontology.json``; an in-process consumer (e.g. langchain) can cache
the :class:`Ontology` in memory. Fetch a fresh one over HTTP with
:meth:`diffbot.Diffbot.dql_fetch_ontology`.
"""

import json
import pathlib
import re
from typing import Any, Dict, List, Optional, Tuple, Union


class Ontology:
    """Queryable view over a parsed Diffbot ontology document.

    The instance holds the parsed document on :attr:`data` and exposes pure
    lookup methods over it. Nothing here performs I/O — construct with already
    parsed data, or use :meth:`from_json` / :meth:`from_path` for convenience.
    """

    def __init__(self, data: Dict[str, Any]):
        self.data = data

    @classmethod
    def from_json(cls, raw: Union[str, bytes]) -> "Ontology":
        """Build from a raw JSON string or bytes (e.g. an HTTP response body)."""
        return cls(json.loads(raw))

    @classmethod
    def from_path(cls, path: Union[str, pathlib.Path]) -> "Ontology":
        """Build from a JSON file on disk."""
        return cls(json.loads(pathlib.Path(path).read_text()))

    def types(self) -> List[str]:
        """All entity type names (e.g. ``Organization``, ``Person``)."""
        return sorted(self.data.get("types", {}).keys())

    def composites(self) -> List[str]:
        """All composite type names (e.g. ``Location``, ``Employment``)."""
        return sorted(self.data.get("composites", {}).keys())

    def enums(self) -> List[str]:
        """All enum type names (e.g. ``Language``, ``Gender``)."""
        return sorted(self.data.get("enums", {}).keys())

    def taxonomies(self) -> List[str]:
        """All taxonomy names (e.g. ``OrganizationCategory``)."""
        return sorted(self.data.get("taxonomies", {}).keys())

    @staticmethod
    def _fields_of(container: Dict[str, Any], type_name: str) -> Dict[str, Any]:
        entry = container.get(type_name)
        if entry is None:
            raise KeyError(f"Unknown name: {type_name}")
        return entry.get("fields", {})

    def fields_for(self, type_name: str) -> Dict[str, Any]:
        """Return the field map of an entity type or composite.

        Auto-routes: ``type_name`` may be an entity type (``Organization``) or a
        composite (``Location``). Raises ``KeyError`` if it is neither.
        """
        types = self.data.get("types", {})
        composites = self.data.get("composites", {})
        if type_name in types:
            return self._fields_of(types, type_name)
        if type_name in composites:
            return self._fields_of(composites, type_name)
        raise KeyError(f"{type_name} is not a known entity type or composite")

    @staticmethod
    def filter_fields(
        fields: Dict[str, Any],
        search: Optional[str],
        include_deprecated: bool = False,
    ) -> List[Tuple[str, Dict[str, Any]]]:
        """Filter a field map by a name regex, dropping deprecated by default."""
        pattern = re.compile(search, re.IGNORECASE) if search else None
        out = []
        for name, meta in fields.items():
            if not include_deprecated and meta.get("isDeprecated"):
                continue
            if pattern and not pattern.search(name):
                continue
            out.append((name, meta))
        return out

    def taxonomy_values(self, name: str, search: Optional[str] = None) -> List[str]:
        """Flatten a taxonomy's values (recursing into children), optionally filtered."""
        tax = self.data.get("taxonomies", {}).get(name)
        if tax is None:
            raise KeyError(f"Unknown taxonomy: {name}")
        pattern = re.compile(search, re.IGNORECASE) if search else None
        out: List[str] = []

        def walk(node: Dict[str, Any]) -> None:
            n = node.get("name")
            if n and (pattern is None or pattern.search(n)):
                out.append(n)
            for child in node.get("children", []) or []:
                walk(child)

        for cat in tax.get("categories", []) or []:
            walk(cat)
        return out

    def enum_values(self, name: str) -> List[str]:
        """Return the allowed values of an enum."""
        enum = self.data.get("enums", {}).get(name)
        if enum is None:
            raise KeyError(f"Unknown enum: {name}")
        return list(enum.get("values", []))

    def find_named(self, search: str) -> List[str]:
        """Fallback search: every ``name`` anywhere in the document matching a regex."""
        pattern = re.compile(search, re.IGNORECASE)
        found = set()

        def walk(node: Any) -> None:
            if isinstance(node, dict):
                n = node.get("name")
                if isinstance(n, str) and pattern.search(n):
                    found.add(n)
                for v in node.values():
                    walk(v)
            elif isinstance(node, list):
                for v in node:
                    walk(v)

        walk(self.data)
        return sorted(found)

    @staticmethod
    def format_field(name: str, meta: Dict[str, Any]) -> str:
        """Render one field as ``<name>: [<type>] [flags...]`` for display."""
        t = meta.get("type", "?")
        if t == "LinkedEntity":
            le = meta.get("leType") or []
            if le:
                t = f"LinkedEntity ({le[0]})"
        flags = []
        if meta.get("isList"):
            flags.append("isList")
        if meta.get("isComposite"):
            flags.append("isComposite")
        if meta.get("isEnum"):
            flags.append("isEnum")
        if meta.get("isDeprecated"):
            flags.append("DEPRECATED")
        suffix = "".join(f" [{f}]" for f in flags)
        return f"{name}: [{t}]{suffix}"