"""
Octopoda × LangChain Integration
==================================
Drop-in memory for LangChain agents and chains.

Setup:
    pip install octopoda langchain langchain-core

Usage with ConversationChain:
    from synrix.integrations.langchain import OctopodaMemory
    from langchain.chains import ConversationChain
    from langchain_openai import ChatOpenAI

    memory = OctopodaMemory(agent_id="support_bot")
    chain = ConversationChain(llm=ChatOpenAI(), memory=memory)
    chain.invoke({"input": "My name is Alice"})
    chain.invoke({"input": "What's my name?"})  # Remembers Alice

Usage with RunnableWithMessageHistory:
    from synrix.integrations.langchain import OctopodaChatHistory

    def get_session_history(session_id):
        return OctopodaChatHistory(agent_id="my_agent", session_id=session_id)

    with_history = RunnableWithMessageHistory(chain, get_session_history)

Usage with AgentExecutor:
    from synrix.integrations.langchain import OctopodaMemory
    memory = OctopodaMemory(agent_id="research_agent")
    agent_executor = AgentExecutor(agent=agent, tools=tools, memory=memory)
"""

from __future__ import annotations

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

# Core Octopoda — always available
from synrix_runtime.api.runtime import AgentRuntime


# ---------------------------------------------------------------------------
# LangChain BaseMemory adapter
# ---------------------------------------------------------------------------

try:
    from langchain_core.memory import BaseMemory
    LANGCHAIN_AVAILABLE = True
except ImportError:
    try:
        from langchain.memory.base import BaseMemory
        LANGCHAIN_AVAILABLE = True
    except ImportError:
        BaseMemory = object
        LANGCHAIN_AVAILABLE = False

try:
    from langchain_core.chat_history import BaseChatMessageHistory
    from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
    CHAT_HISTORY_AVAILABLE = True
except ImportError:
    BaseChatMessageHistory = object
    CHAT_HISTORY_AVAILABLE = False


class OctopodaMemory(BaseMemory):
    """
    Drop-in LangChain memory backed by Octopoda.

    Persists conversation history across restarts. Supports semantic search
    for relevant context retrieval. Works with ConversationChain, AgentExecutor,
    and any LangChain component that accepts a BaseMemory.

    Features:
        - Persistent across process restarts
        - Semantic search for relevant memories
        - Knowledge graph extraction (automatic)
        - Temporal versioning (every update tracked)
        - Works offline, zero API costs
    """

    # Pydantic fields (LangChain uses pydantic for BaseMemory)
    agent_id: str = "langchain_agent"
    memory_key: str = "history"
    input_key: str = "input"
    output_key: str = "output"
    human_prefix: str = "Human"
    ai_prefix: str = "AI"
    return_messages: bool = False
    k: int = 50  # Max conversation turns to return

    # Internal — not a pydantic field
    _runtime: Any = None

    class Config:
        arbitrary_types_allowed = True
        underscore_attrs_are_private = True

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._runtime = AgentRuntime(self.agent_id, agent_type="langchain")

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

    def _get_conversation_key(self) -> str:
        return "conversation_history"

    def load_memory_variables(self, inputs: Dict[str, Any] = None) -> Dict[str, Any]:
        """Load conversation history from Octopoda."""
        result = self._runtime.recall(self._get_conversation_key())

        if not result.found or not result.value:
            if self.return_messages:
                return {self.memory_key: []}
            return {self.memory_key: ""}

        history = result.value
        if isinstance(history, str):
            try:
                history = json.loads(history)
            except (json.JSONDecodeError, TypeError):
                if self.return_messages:
                    return {self.memory_key: []}
                return {self.memory_key: history}

        if not isinstance(history, list):
            if self.return_messages:
                return {self.memory_key: []}
            return {self.memory_key: str(history)}

        # Limit to last k turns
        turns = history[-self.k:] if len(history) > self.k else history

        if self.return_messages:
            messages = []
            for turn in turns:
                if isinstance(turn, dict):
                    if turn.get("role") == "human":
                        messages.append(HumanMessage(content=turn.get("content", "")))
                    elif turn.get("role") == "ai":
                        messages.append(AIMessage(content=turn.get("content", "")))
                    elif turn.get("role") == "system":
                        messages.append(SystemMessage(content=turn.get("content", "")))
            return {self.memory_key: messages}

        # Return as formatted string
        lines = []
        for turn in turns:
            if isinstance(turn, dict):
                role = turn.get("role", "unknown")
                content = turn.get("content", "")
                if role == "human":
                    lines.append(f"{self.human_prefix}: {content}")
                elif role == "ai":
                    lines.append(f"{self.ai_prefix}: {content}")
        return {self.memory_key: "\n".join(lines)}

    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
        """Save a conversation turn to Octopoda."""
        # Load existing history
        result = self._runtime.recall(self._get_conversation_key())
        history = []
        if result.found and result.value:
            val = result.value
            if isinstance(val, str):
                try:
                    val = json.loads(val)
                except (json.JSONDecodeError, TypeError):
                    val = []
            if isinstance(val, list):
                history = val

        # Extract input/output
        human_input = inputs.get(self.input_key, "")
        ai_output = outputs.get(self.output_key, "")
        if not ai_output and isinstance(outputs, dict):
            # Try common output keys
            for key in ("output", "text", "response", "answer"):
                if key in outputs:
                    ai_output = outputs[key]
                    break

        # Append new turn
        history.append({"role": "human", "content": str(human_input), "timestamp": time.time()})
        history.append({"role": "ai", "content": str(ai_output), "timestamp": time.time()})

        # Save back (with temporal versioning)
        self._runtime.remember(self._get_conversation_key(), json.dumps(history))

        # Also store each exchange as a separate searchable memory
        turn_num = len(history) // 2
        self._runtime.remember(
            f"turn_{turn_num}",
            f"{self.human_prefix}: {human_input}\n{self.ai_prefix}: {ai_output}",
        )

    def clear(self) -> None:
        """Clear conversation history."""
        self._runtime.remember(self._get_conversation_key(), json.dumps([]))

    # ----- Bonus: semantic recall for RAG-style memory -----

    def search(self, query: str, limit: int = 5) -> List[Dict]:
        """
        Search past conversations by meaning (semantic search).

        Usage:
            results = memory.search("What did the user say about pricing?")
            for r in results:
                print(r["value"], r["score"])
        """
        result = self._runtime.recall_similar(query, limit=limit)
        return result.items


# ---------------------------------------------------------------------------
# LangChain ChatMessageHistory adapter
# ---------------------------------------------------------------------------

class OctopodaChatHistory(BaseChatMessageHistory):
    """
    Octopoda-backed ChatMessageHistory for RunnableWithMessageHistory.

    Usage:
        from synrix.integrations.langchain import OctopodaChatHistory

        def get_session_history(session_id: str):
            return OctopodaChatHistory(agent_id="my_agent", session_id=session_id)

        chain_with_history = RunnableWithMessageHistory(
            runnable, get_session_history
        )
    """

    def __init__(self, agent_id: str = "langchain_agent", session_id: str = "default"):
        self.agent_id = agent_id
        self.session_id = session_id
        self._runtime = AgentRuntime(agent_id, agent_type="langchain")

    @property
    def messages(self) -> List[BaseMessage]:
        """Retrieve all messages from Octopoda."""
        if not CHAT_HISTORY_AVAILABLE:
            return []

        key = f"chat_history:{self.session_id}"
        result = self._runtime.recall(key)

        if not result.found or not result.value:
            return []

        data = result.value
        if isinstance(data, str):
            try:
                data = json.loads(data)
            except (json.JSONDecodeError, TypeError):
                return []

        if not isinstance(data, list):
            return []

        messages = []
        for msg in data:
            if isinstance(msg, dict):
                content = msg.get("content", "")
                role = msg.get("role", "human")
                if role == "human":
                    messages.append(HumanMessage(content=content))
                elif role == "ai":
                    messages.append(AIMessage(content=content))
                elif role == "system":
                    messages.append(SystemMessage(content=content))
        return messages

    def add_message(self, message: BaseMessage) -> None:
        """Add a message to the history."""
        key = f"chat_history:{self.session_id}"
        result = self._runtime.recall(key)

        history = []
        if result.found and result.value:
            val = result.value
            if isinstance(val, str):
                try:
                    val = json.loads(val)
                except (json.JSONDecodeError, TypeError):
                    val = []
            if isinstance(val, list):
                history = val

        role = "human"
        if isinstance(message, AIMessage):
            role = "ai"
        elif isinstance(message, SystemMessage):
            role = "system"

        history.append({
            "role": role,
            "content": message.content,
            "timestamp": time.time(),
        })

        self._runtime.remember(key, json.dumps(history))

    def clear(self) -> None:
        """Clear all messages."""
        key = f"chat_history:{self.session_id}"
        self._runtime.remember(key, json.dumps([]))
