"""
mubit-langchain — LangChain BaseMemory backed by MuBit.

Provides BaseMemory and BaseChatMemory subclasses that route all
memory operations through MuBit's control plane.
"""

from __future__ import annotations

import json
from typing import Any, Dict, List, Optional

from .mubit_client import MubitControlClient, normalize_metadata

try:
    from langchain_core.memory import BaseMemory
except (ImportError, ModuleNotFoundError):
    try:
        from langchain.memory.base import BaseMemory
    except (ImportError, ModuleNotFoundError):
        # langchain_core >=1.2 removed BaseMemory. Provide a pydantic-based
        # fallback so MubitMemory still works as a standalone memory class.
        from pydantic import BaseModel as _BaseModel

        class BaseMemory(_BaseModel):
            """Standalone BaseMemory when langchain BaseMemory is unavailable."""
            class Config:
                arbitrary_types_allowed = True

try:
    from pydantic import Field
except ImportError:
    from pydantic.v1 import Field


class MubitMemory(BaseMemory):
    """LangChain BaseMemory backed by MuBit memory engine.

    Stores and retrieves conversation context through MuBit's control
    plane, providing semantic memory across chain invocations.

    Usage::

        from mubit_langchain import MubitMemory
        from langchain.chains import ConversationChain

        memory = MubitMemory(api_key="mbt_...", session_id="chain-1")
        chain = ConversationChain(llm=llm, memory=memory)
    """

    endpoint: str = Field(default="http://127.0.0.1:3000")
    api_key: str = Field(default="")
    session_id: str = Field(default="default")
    user_id: str = Field(default="")
    agent_id: str = Field(default="langchain")
    memory_key: str = Field(default="history")
    input_key: Optional[str] = Field(default=None)
    output_key: Optional[str] = Field(default=None)
    return_messages: bool = Field(default=False)
    max_token_budget: int = Field(default=2048)
    entry_types: Optional[List[str]] = Field(default=None)
    context_sections: Optional[List[str]] = Field(default=None)

    _client: Optional[MubitControlClient] = None

    class Config:
        arbitrary_types_allowed = True

    def _get_client(self) -> MubitControlClient:
        if self._client is None:
            self._client = MubitControlClient(
                endpoint=self.endpoint.rstrip("/"),
                api_key=self.api_key,
            )
        return self._client

    @property
    def memory_variables(self) -> List[str]:
        """Return the list of memory variable keys."""
        return [self.memory_key]

    def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Fetch relevant context from MuBit based on the input.

        Args:
            inputs: The chain inputs dict. Uses the first string value
                   as the query for context retrieval.

        Returns:
            Dict with memory_key mapped to the context string.
        """
        query = self._extract_query(inputs)
        if not query:
            return {self.memory_key: "" if not self.return_messages else []}

        client = self._get_client()
        result = client.get_context(
            session_id=self.session_id,
            user_id=self.user_id or None,
            query=query,
            agent_id=self.agent_id,
            max_token_budget=self.max_token_budget,
            entry_types=self.entry_types,
            sections=self.context_sections,
        )

        context = result.get("context_block", "")

        if self.return_messages:
            try:
                from langchain_core.messages import SystemMessage
                return {self.memory_key: [SystemMessage(content=context)] if context else []}
            except ImportError:
                from langchain.schema import SystemMessage
                return {self.memory_key: [SystemMessage(content=context)] if context else []}

        return {self.memory_key: context}

    def save_context(
        self, inputs: Dict[str, Any], outputs: Dict[str, str]
    ) -> None:
        """Store the input/output pair in MuBit.

        Args:
            inputs: The chain inputs.
            outputs: The chain outputs.
        """
        input_text = self._extract_query(inputs)
        output_text = self._extract_output(outputs)

        if not input_text and not output_text:
            return

        client = self._get_client()

        # Ingest the interaction as a trace
        text = f"User: {input_text}\nAssistant: {output_text}"
        client.remember(
            session_id=self.session_id,
            user_id=self.user_id or None,
            text=text,
            agent_id=self.agent_id,
            intent="trace",
            metadata={"source": "langchain"},
        )

    def clear(self) -> None:
        """Clear all memories for the current session."""
        client = self._get_client()
        client.delete_run(session_id=self.session_id)

    def _extract_query(self, inputs: Dict[str, Any]) -> str:
        """Extract query string from inputs dict."""
        if self.input_key:
            return str(inputs.get(self.input_key, ""))
        # Find first string value
        for key, value in inputs.items():
            if isinstance(value, str) and value.strip():
                return value
        return ""

    def _extract_output(self, outputs: Dict[str, str]) -> str:
        """Extract output string from outputs dict."""
        if self.output_key:
            return str(outputs.get(self.output_key, ""))
        for key, value in outputs.items():
            if isinstance(value, str) and value.strip():
                return value
        return ""


class MubitChatMemory(MubitMemory):
    """Chat-oriented LangChain memory backed by MuBit.

    Extends MubitMemory with chat-specific features: stores each
    message individually and returns messages instead of raw text.

    Usage::

        from mubit_langchain import MubitChatMemory

        memory = MubitChatMemory(api_key="mbt_...", session_id="chat-1")
    """

    return_messages: bool = Field(default=True)

    def save_context(
        self, inputs: Dict[str, Any], outputs: Dict[str, str]
    ) -> None:
        """Store individual messages in MuBit."""
        client = self._get_client()

        input_text = self._extract_query(inputs)
        output_text = self._extract_output(outputs)

        if input_text:
            client.remember(
                session_id=self.session_id,
                user_id=self.user_id or None,
                text=input_text,
                agent_id=self.agent_id,
                intent="context",
                metadata={"role": "user", "source": "langchain"},
            )

        if output_text:
            client.remember(
                session_id=self.session_id,
                user_id=self.user_id or None,
                text=output_text,
                agent_id=self.agent_id,
                intent="trace",
                metadata={"role": "assistant", "source": "langchain"},
            )
