"""LangChain integration - NeurovaultMemory with knowledge graph context."""

from __future__ import annotations

import time
from typing import Any, Optional

try:
    from langchain_core.memory import BaseMemory
except ImportError as exc:  # pragma: no cover
    raise ImportError("Install: pip install neurovault[langchain]") from exc

from neurovault.engine import NeurovaultEngine


class NeurovaultMemory(BaseMemory):
    """LangChain memory backed by NEUROVAULT's cognitive engine."""

    vault_name: str = "default"
    memory_key: str = "history"
    input_key: str = "input"
    output_key: str = "output"
    max_memories: int = 10
    include_entities: bool = True
    include_predictions: bool = True
    auto_save: bool = True

    _engine: Optional[NeurovaultEngine] = None

    class Config:
        arbitrary_types_allowed = True
        underscore_attrs_are_private = True

    @property
    def engine(self) -> NeurovaultEngine:
        if self._engine is None:
            self._engine = NeurovaultEngine.open(self.vault_name)
        return self._engine

    @property
    def memory_variables(self) -> list[str]:
        return [self.memory_key]

    def load_memory_variables(self, inputs: dict[str, Any]) -> dict[str, str]:
        query = inputs.get(self.input_key, "")
        if not query:
            return {self.memory_key: ""}

        parts = []

        results = self.engine.recall(query, limit=self.max_memories, mode="multi")
        if results:
            parts.append("Relevant memories:")
            for result in results:
                entity_text = ""
                if self.include_entities and result.entities:
                    entity_text = f" [entities: {', '.join(e.name for e in result.entities[:3])}]"
                parts.append(f"  - {result.memory.content}{entity_text}")

        if self.include_predictions:
            self.engine.update_context("langchain_input", query[:120])
            predictions = self.engine.predict()
            if predictions:
                parts.append("\nPredictive context:")
                for pred in predictions[:3]:
                    parts.append(f"  - {pred.memory.content}")

        return {self.memory_key: "\n".join(parts) if parts else "No relevant context."}

    def save_context(self, inputs: dict[str, Any], outputs: dict[str, str]) -> None:
        if not self.auto_save:
            return

        user_input = inputs.get(self.input_key, "")
        ai_output = outputs.get(self.output_key, "")

        if user_input:
            self.engine.ingest(f"User: {user_input}", source="user", importance=0.5)
        if ai_output and len(ai_output) > 20:
            self.engine.ingest(f"Assistant: {ai_output[:500]}", source="agent", importance=0.4)

    def clear(self) -> None:
        new_name = f"cleared-{int(time.time())}"
        self.engine._bridge.repo.persona.create(new_name)
        self.engine._bridge.repo.persona.use(new_name)
