"""Concept hierarchy and ontology engine."""

from __future__ import annotations

import hashlib
import time
from collections import deque
from typing import Any

from neurovault.storage import NeurovaultStorage


class OntologyEngine:
    """Manages IS-A, PART-OF, and HAS-PROPERTY relationships."""

    DEFAULT_TAXONOMY = {
        "TECHNOLOGY": {
            "Programming Language": [
                "python",
                "javascript",
                "typescript",
                "rust",
                "go",
                "java",
                "ruby",
                "c++",
                "c#",
                "swift",
                "kotlin",
            ],
            "Database": [
                "postgresql",
                "mysql",
                "sqlite",
                "mongodb",
                "redis",
                "cassandra",
                "dynamodb",
                "elasticsearch",
            ],
            "Framework": [
                "react",
                "vue",
                "angular",
                "svelte",
                "django",
                "flask",
                "fastapi",
                "next.js",
                "express",
            ],
            "Cloud Platform": ["aws", "gcp", "azure"],
            "AI/ML Tool": [
                "pytorch",
                "tensorflow",
                "langchain",
                "ollama",
                "huggingface",
                "openai",
                "anthropic",
            ],
            "DevOps Tool": ["docker", "kubernetes", "terraform", "github", "gitlab", "jenkins"],
        },
    }

    def __init__(self, storage: NeurovaultStorage):
        self._storage = storage

    def bootstrap_taxonomy(self) -> int:
        """Load default IS-A taxonomy into the knowledge graph."""
        count = 0
        now = int(time.time() * 1000)

        for root_type, categories in self.DEFAULT_TAXONOMY.items():
            for category_name, children in categories.items():
                category = self._storage.upsert_entity(
                    name=category_name,
                    entity_type=root_type,
                    confidence=1.0,
                    seen_at=now,
                )

                for child_name in children:
                    child = self._storage.upsert_entity(
                        name=child_name,
                        entity_type=root_type,
                        confidence=1.0,
                        seen_at=now,
                    )
                    self._add_hierarchy(child.id, category.id, "IS-A", confidence=0.95, source="bootstrap")
                    count += 1

        return count

    def infer_hierarchy(self, entity_name: str, entity_type: str) -> list[dict[str, Any]]:
        """Infer IS-A relationships for an entity using taxonomy lookup."""
        del entity_type
        inferred: list[dict[str, Any]] = []
        name_lower = entity_name.lower()

        for _root_type, categories in self.DEFAULT_TAXONOMY.items():
            for category_name, children in categories.items():
                if name_lower in [c.lower() for c in children]:
                    inferred.append(
                        {
                            "child": entity_name,
                            "parent": category_name,
                            "relation": "IS-A",
                            "confidence": 0.9,
                        }
                    )

        return inferred

    def get_ancestors(self, entity_id: bytes, max_depth: int = 5) -> list[dict[str, Any]]:
        """Get all ancestors of an entity via IS-A chain."""
        visited = set()
        queue = deque([(entity_id, 0)])
        ancestors: list[dict[str, Any]] = []

        while queue:
            eid, depth = queue.popleft()
            if eid in visited or depth > max_depth:
                continue
            visited.add(eid)

            parents = self._storage.get_hierarchy_parents(eid, "IS-A")
            for parent_id, rel_type, confidence in parents:
                try:
                    parent = self._storage.get_entity(parent_id)
                    ancestors.append(
                        {
                            "entity": parent,
                            "depth": depth + 1,
                            "relation": rel_type,
                            "confidence": confidence,
                        }
                    )
                    queue.append((parent_id, depth + 1))
                except Exception:
                    continue

        return ancestors

    def get_descendants(self, entity_id: bytes, max_depth: int = 3) -> list[dict[str, Any]]:
        """Get all descendants of an entity via IS-A chain."""
        visited = set()
        queue = deque([(entity_id, 0)])
        descendants: list[dict[str, Any]] = []

        while queue:
            eid, depth = queue.popleft()
            if eid in visited or depth > max_depth:
                continue
            visited.add(eid)

            children = self._storage.get_hierarchy_children(eid, "IS-A")
            for child_id, rel_type, confidence in children:
                try:
                    child = self._storage.get_entity(child_id)
                    descendants.append(
                        {
                            "entity": child,
                            "depth": depth + 1,
                            "relation": rel_type,
                            "confidence": confidence,
                        }
                    )
                    queue.append((child_id, depth + 1))
                except Exception:
                    continue

        return descendants

    def get_siblings(self, entity_id: bytes) -> list[dict[str, Any]]:
        """Get entities sharing the same parent (siblings in IS-A tree)."""
        parents = self._storage.get_hierarchy_parents(entity_id, "IS-A")
        siblings: list[dict[str, Any]] = []
        seen = {entity_id}

        for parent_id, _, _ in parents:
            children = self._storage.get_hierarchy_children(parent_id, "IS-A")
            for child_id, _, confidence in children:
                if child_id in seen:
                    continue
                seen.add(child_id)
                try:
                    child = self._storage.get_entity(child_id)
                    siblings.append(
                        {
                            "entity": child,
                            "shared_parent_id": parent_id,
                            "confidence": confidence,
                        }
                    )
                except Exception:
                    continue

        return siblings

    def expand_query_with_ontology(self, entity_names: list[str]) -> list[str]:
        """Expand entity names with ontology descendants."""
        expanded = list(entity_names)

        for name in entity_names:
            matches = self._storage.find_entities_by_name(name)
            for entity in matches:
                descendants = self.get_descendants(entity.id, max_depth=2)
                for desc in descendants:
                    entity_name = desc["entity"].name
                    if entity_name not in expanded:
                        expanded.append(entity_name)

        return expanded

    def _add_hierarchy(
        self,
        child_id: bytes,
        parent_id: bytes,
        relation: str,
        confidence: float = 0.8,
        source: str = "inferred",
    ) -> None:
        hid = hashlib.sha256(child_id + parent_id + relation.encode("utf-8")).digest()
        now = int(time.time() * 1000)
        self._storage.insert_concept_hierarchy(
            hid,
            child_id,
            parent_id,
            relation,
            confidence,
            source,
            now,
        )
