"""
HieromemMemory - LangChain BaseMemory implementation.

Drop-in replacement for LangChain conversation memory that persists to HIEROMEM.
"""
from typing import Any, Dict, List, Optional
import httpx

try:
    from langchain.memory.chat_memory import BaseChatMemory
    from langchain.schema import BaseMessage, HumanMessage, AIMessage
    LANGCHAIN_AVAILABLE = True
except ImportError:
    LANGCHAIN_AVAILABLE = False
    BaseChatMemory = object
    BaseMessage = object
    HumanMessage = object
    AIMessage = object


class HieromemMemory(BaseChatMemory):  # type: ignore[misc]
    """
    LangChain memory backed by HIEROMEM.

    Usage:
        from hieromem.integrations.langchain import HieromemMemory

        memory = HieromemMemory(
            hieromem_url="http://localhost:8000",
            api_key="sk-...",
            namespace="chat_history"
        )

        chain = ConversationChain(llm=llm, memory=memory)
    """

    hieromem_url: str = "http://localhost:8000"
    api_key: Optional[str] = None
    namespace: str = "default"
    return_messages: bool = True
    memory_key: str = "history"
    input_key: str = "input"
    output_key: str = "output"
    k: int = 10  # Number of recent messages to retrieve

    _client: Optional[httpx.Client] = None

    class Config:
        arbitrary_types_allowed = True
        extra = "allow"

    def __init__(self, **kwargs: Any):
        if not LANGCHAIN_AVAILABLE:
            raise ImportError(
                "LangChain is required for HieromemMemory. "
                "Install with: pip install langchain"
            )
        super().__init__(**kwargs)
        self._client = httpx.Client(
            base_url=self.hieromem_url,
            headers={"Authorization": f"Bearer {self.api_key}"} if self.api_key else {},
            timeout=30.0,
        )

    @property
    def memory_variables(self) -> List[str]:
        """Return memory variables."""
        return [self.memory_key]

    def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """
        Load memory variables from HIEROMEM.

        Retrieves recent conversation history relevant to the input.
        """
        query = inputs.get(self.input_key, "")

        if self._client is None:
            return {self.memory_key: [] if self.return_messages else ""}

        try:
            response = self._client.post(
                "/memory/retrieve",
                json={
                    "query": query,
                    "k": self.k,
                    "namespace": self.namespace,
                    "mode": "graph_hybrid",
                },
            )
            response.raise_for_status()
            episodes = response.json().get("episodes", [])

            if self.return_messages:
                messages = self._episodes_to_messages(episodes)
                return {self.memory_key: messages}
            else:
                # Return as string
                history = self._episodes_to_string(episodes)
                return {self.memory_key: history}

        except httpx.HTTPError:
            # Fallback to empty on error
            return {self.memory_key: [] if self.return_messages else ""}

    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
        """
        Save conversation context to HIEROMEM.

        Persists both human input and AI output as separate episodes.
        """
        human_input = inputs.get(self.input_key, "")
        ai_output = outputs.get(self.output_key, "")

        # Save human message
        if human_input:
            self._insert_episode(
                content=human_input,
                role="human",
            )

        # Save AI message
        if ai_output:
            self._insert_episode(
                content=ai_output,
                role="ai",
            )

    def _insert_episode(self, content: str, role: str) -> None:
        """Insert a single episode to HIEROMEM."""
        if self._client is None:
            return

        try:
            self._client.post(
                "/memory/insert",
                json={
                    "text": content,
                    "namespace": self.namespace,
                    "meta": {
                        "role": role,
                        "type": "conversation",
                    },
                },
            )
        except httpx.HTTPError:
            pass  # Silent fail on insert errors

    def _episodes_to_messages(self, episodes: List[Dict[str, Any]]) -> List[BaseMessage]:
        """Convert HIEROMEM episodes to LangChain messages."""
        messages = []
        for ep in episodes:
            meta = ep.get("meta", {})
            content = meta.get("text", "")
            role = meta.get("role", "human")

            if role == "human":
                messages.append(HumanMessage(content=content))
            else:
                messages.append(AIMessage(content=content))

        return messages

    def _episodes_to_string(self, episodes: List[Dict[str, Any]]) -> str:
        """Convert HIEROMEM episodes to string format."""
        lines = []
        for ep in episodes:
            meta = ep.get("meta", {})
            content = meta.get("text", "")
            role = meta.get("role", "Human")
            lines.append(f"{role.title()}: {content}")
        return "\n".join(lines)

    def clear(self) -> None:
        """Clear memory by deleting all episodes in namespace."""
        if self._client is None:
            return

        try:
            self._client.delete(
                f"/memory/namespace/{self.namespace}",
            )
        except httpx.HTTPError:
            pass

    def __del__(self) -> None:
        """Cleanup HTTP client."""
        if self._client:
            self._client.close()
